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

问题描述

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

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

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

我可以在不改变功能的情况下做到这一点吗？如果没有，您将如何更改功能？

解决方法

在将函数应用于它们的第一个和第二个元素之前，您可以取这些值的笛卡尔积：

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

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

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