将 numpy 操作转换为 tensorflow

Question

I am looking to translate numpy operations to tensorflow.

A function given a 2D ndarray, I want to make all entries that are not the maximum along axis 0 update to the value of 0.

def rows_to_zero(arr: np.ndarray):
    def row_to_zero(row: np.ndarray):
        row[row < row.max()] = 0
        return row
    return np.apply_along_axis(row_to_zero, 0, arr)

In: [[1, 2, 1, 4, 5, 3],
     [2, 4, 5, 0, 1, 3],
     [3, 5, 3, 6, 7, 1]]

Out: [[0, 0, 0, 0, 5, 0],
      [0, 0, 5, 0, 0, 0],
      [0, 0, 0, 0, 7, 0]]

I want to write this same functionality using Tensorflow tensors

Would anyone be able to help with something like this?

Answer 1

You just need tf.reduce_max and tf.where .

import tensorflow as tf

a = tf.constant([[1, 2, 1, 4, 5, 3],
                 [2, 4, 5, 0, 1, 3],
                 [3, 5, 3, 6, 7, 1]],tf.float32)

b = tf.reduce_max(a,axis=1,keepdims=True)
result = tf.where(tf.less(a,b),tf.zeros_like(a),a)

with tf.Session() as sess:
    print(sess.run(result))

# [[0. 0. 0. 0. 5. 0.]
#  [0. 0. 5. 0. 0. 0.]
#  [0. 0. 0. 0. 7. 0.]]