python - 张量流中对张量对象进行特定类型的切片

Question

问题总结，张量流是否支持这种切片然后赋值？

out[tf_a2[y],x[:,None]] = tf_a1[tf_a2[y],x[:,None]]
final = out[:-1]

让我们举个例子，我有一个这样的张量：

tf_a1 = tf.Variable([    [9.968594,  8.655439,  0.,        0.       ],
                         [0.,        8.3356,    0.,        8.8974   ],
                         [0.,        0.,        6.103182,  7.330564 ],
                         [6.609862,  0.,        3.0614321, 0.       ],
                         [9.497023,  0.,        3.8914037, 0.       ],
                         [0.,        8.457685,  8.602337,  0.       ],
                         [0.,        0.,        5.826657,  8.283971 ],
                         [0.,        0.,        0.,        0.       ]])

我有这个：

tf_a2 = tf.constant([[1, 2, 5],
                    [1, 4, 6],
                    [0, 7, 7],
                    [2, 3, 6],
                    [2, 4, 7]])

现在我想保留tf_a1其中 n （这里 n 是 2）的组合（它们的索引）在值中的元素tf_a2。这是什么意思？

例如，在 中tf_a1，在第一列中，具有值的索引是：(0,3,4)。是否有任何行tf_a2包含这两个索引的任意组合：(0,3)、(0,4) 或 (3,4)。实际上，没有这样的行。所以该列中的所有元素都变为零。

第二列的索引tf_a1是 (0,1) (0,5) (1,5)。如您所见，记录 (1,5)tf_a2在第一行中可用。这就是为什么我们将它们保留在tf_a1.

这是正确的numpy代码：

y,x = np.where(np.count_nonzero(a1p[a2], axis=1) >= n)
out = np.zeros_like(tf_a1)
out[tf_a2[y],x[:,None]] = tf_a1[tf_a2[y],x[:,None]]
final = out[:-1]

这是这个 numpy 代码的预期输出（但我在 tensorflow 中需要它）：

[[0.        0.        0.        0.       ]
 [0.        8.3356    0.        8.8974   ]
 [0.        0.        6.103182  7.330564 ]
 [0.        0.        3.0614321 0.       ]
 [0.        0.        3.8914037 0.       ]
 [0.        8.457685  8.602337  0.       ]
 [0.        0.        5.826657  8.283971 ]]

张量流代码应该是这样的：

y, x = tf.where(tf.count_nonzero(tf.gather(tf_a1, tf_a2, axis=0), axis=1) >= n)
out = tf.zeros_like(tf_a1)
out[tf_a2[y],x[:,None]] = tf_a1[tf_a2[y],x[:,None]]
final = out[:-1]

这部分代码tf.gather(tf_a1, tf_a2, axis=0), axis=1)正在做 numpy 之类的切片tf_a1[tf_a2]

更新 1

唯一不起作用的行是：

out[tf_a2[y],x[:,None]] = tf_a1[tf_a2[y],x[:,None]]
final = out[:-1]

知道如何在张量流中实现这一点，张量对象是否支持这种切片？

任何帮助表示赞赏:)

标签： pythonarraystensorflowslice