[FIXED] Match pytorch scatter output in tensorflow

Issue

How can I do the same operation in tensorflow?

tensor = np.random.RandomState(42).uniform(size=(2, 4, 2)).astype(np.float32)
tensor = torch.from_numpy(tensor)
index = tensor.max(-1, keepdim=True)[1]
output = torch.zeros_like(tensor).scatter_(-1, index, 1.0)

expected output:
tensor([[[0., 1.],
         [1., 0.],
         [1., 0.],
         [0., 1.]],

        [[0., 1.],
         [0., 1.],
         [1., 0.],
         [0., 1.]]])

Solution

As always, everything is a bit more complicated with Tensorflow:

import tensorflow as tf
import numpy as np

tensor = np.random.RandomState(42).uniform(size=(2, 4, 2)).astype(np.float32)
tensor = tf.constant(tensor)
_, indices = tf.math.top_k(tensor)
zeros = tf.zeros_like(tensor)

ij = tf.stack(tf.meshgrid(
    tf.range(zeros.shape[0], dtype=tf.int32), 
    tf.range(zeros.shape[1], dtype=tf.int32),
                              indexing='ij'), axis=-1)
gathered_indices = tf.concat([ij, indices], axis=-1)
indices_shape = tf.shape(indices)
values = tf.ones((indices_shape[0], indices_shape[1]))
output = tf.tensor_scatter_nd_update(zeros, gathered_indices, values)

print(output)

tf.Tensor(
[[[0. 1.]
  [1. 0.]
  [1. 0.]
  [0. 1.]]

 [[0. 1.]
  [0. 1.]
  [1. 0.]
  [0. 1.]]], shape=(2, 4, 2), dtype=float32)

Answered By - AloneTogether