CODEWITHSUNDEEP
pythonpytorchvectorizationcartesian-product
Saesun Kim
Saesun Kim

Reputation: 223

Apply vectorised function to Cartesian product of two ranges in PyTorch

I am kind of new to pytorch, and have a very simple question. Let's say we have a scalar function f():

def f(x,y):
    return np.cos(x)+y

What I want to do is use the GPU to generate all pairs of data-points from two ranges x and y. For simple case, take x=y=[0,1,2].

Can I do that without changing the function? If not, how would you change the function?

Upvotes: 0

Views: 722

Answers (1)

iacob
iacob

Reputation: 24181

You can take the Cartesian product of the values before applying your function to their first and second elements:

x = y = torch.tensor([0,1,2])

pairs = torch.cartesian_prod(x,y)
# tensor([[0, 0], [0, 1], [0, 2], [1, 0], ​[1, 1], [1, 2], [2, 0], [2, 1], [2, 2]])

x_, y_ = pairs[:,0], pairs[:,1]
f(x_,y_)

Upvotes: 1

Related Questions