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

Apply vectorised function to Cartesian product of two ranges in PyTorch

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

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

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

我可以在不更改功能的情况下这样做吗?如果没有,您将如何更改功能?

您可以在将函数应用到它们的第一个和第二个元素之前对这些值进行笛卡尔积:

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