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)
Answers:
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)>
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)
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)>