tensorflow2.0 - Tensorflow: binary mask of max values along tensor axis

Question

If I have a N-dimensional tensor, I would like to create another tensor (with the same shape) of values 0 and 1, where 1 is in the same position as the maximum element in original tensor across some dimension.
One constraint I have is that I want to get only the first maximum element along that axis, in case there are duplicates.

For simplification, I will use fewer dimensions.

>>> x = tf.constant([[7, 2, 3], 
                     [5, 0, 1], 
                     [3, 8, 2]], dtype=tf.float32)

>>> tf.reduce_max(x, axis=-1)
tf.Tensor([7. 5. 8.], shape=(3,), dtype=float32)

What I want is:

tf.Tensor([1. 0. 0.], 
          [1. 0. 0.],
          [0. 1. 0.], shape=(3,3), dtype=float32)
          

What I've tried (and realized was wrong):

>>> tf.cast(tf.equal(x, tf.reduce_max(x, axis=-1, keepdims=True)), dtype=tf.float32)

# works fine when there are no duplicates
tf.Tensor([[1. 0. 0.]
           [1. 0. 0.]
           [0. 1. 0.]], shape=(3, 3), dtype=float32)


>>> y = tf.zeros([3,3])
>>> tf.cast(tf.equal(y, tf.reduce_max(y, axis=-1, keepdims=True)), dtype=tf.float32)

# fails when there are multiple identical values across dimension
tf.Tensor([[1. 1. 1.]
           [1. 1. 1.]
           [1. 1. 1.]], shape=(3, 3), dtype=float32)

Edit: Solved

tf.cast(tf.equal(tf.argsort(tf.argsort(x, 1, direction='DESCENDING'), 1), 0), tf.float32)

question from:https://stackoverflow.com/questions/66048913/tensorflow-binary-mask-of-max-values-along-tensor-axis

Answer 1

You could use double tf.argsort() to get rank order of elements along axis 1 and get max rank. This gives the last instance of the max value as the top rank. Let's take an example with duplicate elements -

x = tf.constant([[7, 2, 3],  #max is 7
                 [5, 0, 5],  #max is 5 but duplicate in same row
                 [7, 8, 7]]) #max is 8 but shares 7 with first row too

tf.cast(tf.equal(tf.argsort(tf.argsort(x, 1), 1), x.shape[0]-1), tf.int64)

<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 0, 0],
       [0, 0, 1],
       [0, 1, 0]], dtype=int32)>