How to index and assign to a tensor in tensorflow?

Question

I have a tensor as follows and a numpy 2D array

k = 1
mat = np.array([[1,2],[3,4],[5,6]])
for row in mat:
    values_zero, indices_zero = tf.nn.top_k(row, len(row) - k)
    row[indices_zero] = 0  #????

I want to assign the elements in that row to be zero at those indices. However I can't index a tensor and assign to it as well. I have tried using the tf.gather function but how can I do an assignment? I want to keep it as a tensor and then run it in a session at the end if that is possible.

Answer 1

I guess you are trying to mask the maximum in each row to zero? If so, I would do it like this. The idea is to create the tensor by construction rather than assignment.

import numpy as np
import tensorflow as tf

mat = np.array([[1, 2], [3, 4], [5, 6]])

# All tensorflow from here
tmat = tf.convert_to_tensor(mat)

# Get index of maximum
max_inds = tf.argmax(mat, axis=1)

# Create an array of column indices in each row
shape = tmat.get_shape()
inds = tf.range(0, shape[1], dtype=max_inds.dtype)[None, :]

# Create boolean mask of maximums
bmask = tf.equal(inds, max_inds[:, None])

# Convert boolean mask to ones and zeros
imask = tf.where(bmask, tf.zeros_like(tmat), tf.ones_like(tmat))

# Create new tensor that is masked with maximums set to zer0
newmat = tmat * imask

with tf.Session() as sess:
    print(newmat.eval())

which outputs

[[1 0]
 [3 0]
 [5 0]]

Answer 2

One way to do this is by advanced indexing:

In [87]: k = 1
In [88]: mat = np.array([[1,2],[3,4],[5,6]])

# `sess` is tf.InteractiveSession()
In [89]: vals, idxs = sess.run(tf.nn.top_k(mat, k=1))

In [90]: idxs
Out[90]: 
array([[1],
       [1],
       [1]], dtype=int32)

In [91]: mat[:, np.squeeze(idxs)[0]] = 0

In [92]: mat
Out[92]: 
array([[1, 0],
       [3, 0],
       [5, 0]])