Convert tensor of (row, column) coordinates to boolean mask in TensorFlow-CodePudding

I have an array of 2D coordinates, from which I need to obtain a boolean mask with a known shape, where elements whose index is in the coordinates array is True.

For example, if I had an indices tensor which contains [[0, 0], [1, 1], [2, 2], [0, 1], [1, 2]] and a given shape of (5, 5) I need to get a matrix that is like so:

[[ True,  True, False, False, False],
 [False,  True,  True, False, False],
 [False, False,  True, False, False],
 [False, False, False, False, False],
 [False, False, False, False, False]]

In Numpy I'd do it like so:

idx = np.array([[0, 0], [1, 1], [2, 2], [0, 1], [1, 2]])
bool_mat = np.zeros(shape=(5, 5), dtype=np.bool)
bool_mat[idx[:, 0], idx[:, 1]] = True

However, in TensorFlow you can't assign tensor items like this.

How can I express an equivalent computation in TensorFlow?

CodePudding user response：

You can use tf.tensor_scatter_nd_update(tensor, indices, updates):

import tensorflow as tf

indices = tf.constant([[0, 0], [1, 1], [2, 2], [0, 1], [1, 2]])
bool_mat = tf.zeros(shape=(5, 5), dtype=tf.bool)

bool_mat = tf.tensor_scatter_nd_update(bool_mat, indices, tf.repeat([True], repeats=tf.shape(indices)[0]))
print(bool_mat)

tf.Tensor(
[[ True  True False False False]
 [False  True  True False False]
 [False False  True False False]
 [False False False False False]
 [False False False False False]], shape=(5, 5), dtype=bool)