如何将函数应用于二维张量的每个元素？

Question

我正在寻找一种将tensorflow op应用于2d张量的每个元素的方法。

input = tf.Variable([[1.0, 2.0], [3.0, 4.0]])
# final result should look like: 
# [[myCustomOp(1.0), myCustomOp(2.0)]), [myCustomOp(3.0), myCustomOp(4.0)]]

Answer 1

以下代码有效：

  def myOp(t):
    return t+1
  shape = tf.shape(input)
  elems = tf.reshape(input, [-1])
  res = tf.map_fn(fn=lambda t: myOp(t), elems=elems)
  res = tf.reshape(res, shape)