CC @fmassa

@zasdfgbnm we'd love to have this and I am working on reviewing your code. I had a question about the behavior differences between this implementation and itertools -- please let me know if the semantics of this implementation are intended.

@zou3519 I've changed the behavior of those functions,it looks much more natural now. Please take a look.

Thank you @zasdfgbnm, I'll try to take a look soon.

I have a general question: for cartesian_prod, itertools.product does the following:

In [37]: x = torch.tensor([1, 2])

In [38]: y = torch.tensor([1, 2, 3])

In [39]: list(itertools.product(x, y))
Out[39]:
[(tensor(1), tensor(1)),
 (tensor(1), tensor(2)),
 (tensor(1), tensor(3)),
 (tensor(2), tensor(1)),
 (tensor(2), tensor(2)),
 (tensor(2), tensor(3))]

This is pretty close to what cartesian_prod does: is it a significant usability help and/or speedup that torch.cartesian_product turns all of this into one tensor?

@zou3519 Answer to your question:

1. The speed is faster. The following program:

import torch
import itertools
from tqdm import tqdm

a = torch.rand(100)
b = torch.rand(100)

for _ in tqdm(range(10000)):
	torch.cartesian_prod([a, b])

for _ in tqdm(range(10000)):
	torch.tensor(list(itertools.product(a, b)))

On my computer, the first would be 7806.71it/s, but the second is only 65.69it/s.

2. For convenience, I'm assuming, when we need to do cartesian product or combinations, the input tensor is usually a index tensor, and the result will further be used in something like index_select, index, gather, etc. It would be more convenient if we can directly have a tensor, instead of list of tensors.

3. Is itertools supported in JIT??

@zou3519 I've made changes according to your advice.

Fix the lint please: ./test/test_torch.py:9264:5: E303 too many blank lines (2)

Otherwise, this PR looks good (see my one comment), sorry that it took us so long to get back to you, @zasdfgbnm

@zou3519 Another question: for cartesian_prod, if the input is a single scalar what should the output be? A scalar? A tensor of shape (1,)? I have a feeling that scalar input is confusing and might not very useful, so I turn off support for all scalar inputs. What do you think?

Also, if you think cartesian_prod should support scalar input, then how about combinations? Currently it only support 1d vector input.

I also think if the input has shape (N,) the output should also have shape (N,) instead of (N,1)

@zasdfgbnm, that is a good question. I'm not sure what the right answer here, especially because numpy doesn't have an equivalent to these functions. Let me do some digging around and get back to you.

Edit: I thought about it a little more. I agree that the scalar input is confusing; it would be good to turn off the support for it. In my mind cartesian product should be an operation between two "lists" or 1d tensors. Same with combinations; if someone is trying to take combinations of a scalar tensor I think they may be doing something wrong.

Which operator are you talking about for

I also think if the input has shape (N,) the output should also have shape (N,) instead of (N,1)
?

For cartesian product I agree: if we have an input of shape (N,), the output should have shape (N,) because there was only one input in the list.

@zou3519

Which operator are you talking about for

I also think if the input has shape (N,) the output should also have shape (N,) instead of (N,1)

Here I mean nothing other than simply agreeing your comment at: #9393 (comment)

So now I think we are on complete agreement about this now.

Awesome, @zasdfgbnm. Let me take one last pass through this PR :)

@zasdfgbnm sorry, I forgot about this PR. I'll take another look through it today (I don't remember if there were any unresolved points that we brought over this PR before?)

@zasdfgbnm you can ignore the last build failure, it is unrelated to your code

@zou3519 I've just made changes according to your new comments. I don't think we have unresolved issue last time, but I forgot to answer your comment at #9393 (comment) , which I can answer now:

JIT support is already working:

(pt) ^^/pytorch >>> cat quicktest.py                                                                                            : (*itertools+15746) 15:06:10 
import torch
@torch.jit.script
def f(x):
    return torch.combinations(x)

print(f.graph)
(pt) ^^/pytorch >>> python quicktest.py                                                                                         : (*itertools+15746) 15:06:24 
graph(%x : Tensor) {
  %2 : bool = prim::Constant[value=0]()
  %1 : int = prim::Constant[value=2]()
  %3 : Tensor = aten::combinations(%x, %1, %2)
  return (%3);
}