Python 如何对tensorflow中的张量进行索引和赋值？

我有一个张量，如下所示，还有一个numpy 2D数组

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  #????

---

我想在这些索引中将该行中的元素赋值为零。然而，我不能索引一个张量，也不能给它赋值。我已经尝试过使用tf.gather函数，但是我如何才能完成作业？我想把它作为一个张量，如果可能的话，在最后的会话中运行它。

我想你是想把每行的最大值设为零？如果是这样，我会这样做。其思想是通过构造而不是赋值来创建张量

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

哪个输出

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

---

一种方法是通过高级索引：

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