张量流中是否有一个函数可以进行索引函数的转换？

Question

我正在寻找（但完全找不到）张量流中的一个函数，它允许我在张量上做一个“映射”。

地图

首先，我什至不确定是否有“地图”功能？通过这种方式，可以让我将给定的 f(x) 应用于甚至进入张量。例如我想要这样的东西

def f(x):
  return x**2

Y = tf.Variable(np.array([[1.0, 2.0],
                          [3.0, 4.0]])

Y = tf.map_function(X, f)

生成（显然是在会话中适当运行之后）具有值的张量

Y = [[1.0, 4.0],
     [9.0, 16.0]]

这是否存在（对于一般f - 我意识到tf.nn.relu 和tf.nn.sigmoid? On one hand, it seems like it should, sincemap` 是一个非常基本的操作。另一方面，它会涉及获取提供的python 函数并以某种方式将其转换为在 GPU 上执行，这听起来可能是不可能的。

我在这里是在求一根棍子上的月亮吗？

**地图*

如果存在这样的功能，是否有允许我使用索引感知f 的版本？例如

def f(x, i):
  if (i != [0, 0]):
    k2 = np.sum([x**2 for x in i])
  else:
    k2 = 1.0 # To avoid division by zero
  return (x / k2)

Y = tf.Variable(np.ones(shape=(2,3)))

Y = tf.mapi_function(X, f)

生产

Y = [[1.0, 1.0, 0.25],
     [1.0, 0.5, 0.2]]

如果这样的功能不存在，我是否可以（对于固定的f）通过从（稍微修改的）源构建 tensorflow 来添加它们？

为什么我需要这样的功能

我问这个的原因是我正在尝试使用 tensorflow 对 PDE 进行数值积分。作为其中的一部分，我需要计算拉普拉斯算子 (d^2/dx^2 + d^2/dy^2 + d^2/dz^2) u(x,y,z)。在域 u(k_X, k_y, k_z) 的傅里叶变换表示中，这涉及除以 k_x^2 + k_y^2 + k_z^2。

我可以预先计算一个反平方波数值的张量，然后按元素进行乘法运算。但这会占用大量内存。我怀疑从内存中加载这些值也会更慢。

Answer 1

在您希望单独映射到每个x,y,z 坐标的具体示例中，您可以使用tf.split() 和tf.stack() 轻松完成此操作。也就是说，我假设您有一个大小为[n,m,...,3] 的输入张量（称为K）；也就是说，最后一个维度索引x,y,z 坐标。如果是这样，则使用tf.split() 将K 分解为Kx,Ky,Kz。然后应用您的地图操作（我通常使用tf.map_fn() 来实现此目的），然后最后将它们与tf.stack() 堆叠在一起。

如果我正确理解了应该这样做的设置。如果不是，请提供一个使问题具体化的最小工作示例；否则我们最多只能猜测解决方案。