python - tensorflow 中是否有一个函数可以对索引函数进行转换？

Question

我正在寻找(但完全无法找到) tensorflow 中的一个函数，它允许我在张量上进行“映射”。

map

首先，我什至不确定是否有“ map ”功能？这意味着我可以将给定的 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)

生成(显然在 session 中适当运行之后)具有值的张量

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

这是否存在(对于一般的f - 我意识到tf.nn.relu和tf.nn.sigmoid？一方面，看起来像它应该，因为 map` 是一个非常基本的操作。另一方面，它会涉及获取提供的 python 函数并以某种方式将其转换为在 GPU 上执行，这听起来像是不可能的事情.

我在这里要的是棒上月亮吗？

**mapi*

如果存在这样的函数，是否有一个版本允许我使用索引感知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]]

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

为什么我需要这样的功能

我问这个问题的原因是我正在尝试使用 tensorflow 对偏微分方程进行数值积分。作为其中的一部分，我需要计算拉普拉斯 (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。然后应用 map 操作(我通常使用 tf.map_fn() 来实现此目的)，最后使用 tf.stack() 将内容堆叠在一起。

如果我正确理解了设置，应该可以做到。如果没有，请提供一个最小的工作示例，使问题具体化；否则我们充其量只是猜测解决方案。