locations of maximum values of tensorflow tensor across an axis

Question

I have a rank 3 tensor and another empty tensor of same shape. I am trying to find the maximum values across 3rd(Z) axis for all X,Y locations and insert 1 into the corresponding location in the empty tensor. I have achieved this in numpy in the following way

a = np.random.rand(5,5,3)>=0.5
empty_tensor = np.zeros((5,5,3))
max_z_indices = a.argmax(axis=-1)
empty_tensor[np.arange(a.shape[0])[:,None],np.arange(a.shape[1]),max_z_indices] = 1

In tensorflow I have

a_tf = tf.Variable(a)
empty_tensor_tf = tf.Variable(np.zeros((5,5,3)))
max_z_indices = sess.run(tf.argmax(a_tf,axis=-1))

I know that I can explicitly write the X,Y,Z indices of max values along 3rd dimension of tensor a_tf and use tf.scatter_nd_update to update empty_tensor_tf but I was hoping to find a better way (broadcasting) as in the last line of numpy code.

zihaozhihao zihaozhihao · Accepted Answer · 2019-10-10T23:55:37

You can use tf.reduce_max to get the max value for each z-index, then use tf.where to convert it to 1 or 0 according to the cond.

import tensorflow as tf

# tf_a is (5,5,3) tensor
max_val = tf.reduce_max(tf_a, axis=-1,keepdims=True)
cond = tf.equal(tf_a, max_val)
res = tf.where(cond, tf.ones_like(tf_a), tf.zeros_like(tf_a))

locations of maximum values of tensorflow tensor across an axis

1 Answers