如果我有一个N维Tensor，我想创建另一个值为0和1的Tensor（具有相同的形状），其中1在某个维度上与原始Tensor中的最大元素处于相同的位置。
我的一个约束是，我只想得到第一个最大元素沿该轴，以防有重复。
为了简化，我将使用较少的尺寸。

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

我想要的是：

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

我尝试过（并意识到是错误的）：

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

• • 编辑：已解决**

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