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

Question

我正在寻找一种方法，将张量流运算应用于二维张量的每个元素，例如

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)