将矢量化函数应用于PyTorch中两个范围的笛卡尔乘积

Question

我是pytorch的新手，有一个非常简单的问题。假设我们有一个标量函数f():

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

我想做的是使用GPU生成来自两个范围x和y的所有数据点对。对于简单的情况，取x=y=[0,1,2]。

我可以在不更改函数的情况下执行此操作吗？如果没有，您将如何更改函数？

Answer 1

在将函数应用于值的第一个和第二个元素之前，可以使用值的笛卡尔乘积：

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_)