使用 tf.where() 通过 2D 条件选择 3D 张量，并用键和值替换 2D 索引中的元素

Question

标题中有两个问题。我对这两个问题都感到困惑，因为tensorflow是一种静态编程语言(我真的想回到pytorch或chainer(。

我举了两个例子。请用tensorflow代码或提供相关功能链接回答我。

1( tf.where((

data0 = tf.zeros([2, 3, 4], dtype = tf.float32)
data1 = tf.ones([2, 3, 4], dtype = tf.float32)
cond = tf.constant([[0, 1, 1], [1, 0, 0]])
# cond.shape == (2, 3)
# tf.where() works for 1d condition with 2d data, 
# but not for 2d indices with 3d tensor
# currently, what I am doing is:
#    cond = tf.stack([cond] * 4, 2)
data = tf.where(cond > 0, data1, data0)
# data should be [[0., 1., 1.], [1., 0., 0.]]

(我不知道如何广播秒到三维张量(

2( 二维张量中的变化单元

# all dtype == tf.int64
t2d = tf.Variable([[0, 1, 2], [3, 4, 5]])
k, v = tf.constant([[0, 2], [1, 0]]), tf.constant([-2, -3])
# TODO: change values at positions k to v
# I cannot do [t2d.copy()[i] = j for i, j in k, v]
t3d == [[[0, 1, -2], [3, 4, 5]],
[[0, 1, 2], [-3, 4, 5]]]

事先非常感谢。XD-

Answer 1

这是两个完全不同的问题，它们可能应该这样发布，但无论如何。

1(

是的，您需要手动将所有输入广播到[tf.where](https://www.tensorflow.org/api_docs/python/tf/where]如果它们不同的话。就价值而言，它有一个(旧的(公开问题，但到目前为止，它还没有得到实施。你可以像你建议的那样使用tf.stack，尽管tf.tile可能更明显(并且可能节省内存，尽管我不确定它是如何实现的(：

cond = tf.tile(tf.expand_dims(cond, -1), (1, 1, 4))

或者简单地使用tf.broadcast_to:

cond = tf.broadcast_to(tf.expand_dims(cond, -1), tf.shape(data1))

2(

这是一种方法：

import tensorflow as tf
t2d = tf.constant([[0, 1, 2], [3, 4, 5]])
k, v = tf.constant([[0, 2], [1, 0]]), tf.constant([-2, -3])
# Tile t2d
n = tf.shape(k)[0]
t2d_tile = tf.tile(tf.expand_dims(t2d, 0), (n, 1, 1))
# Add aditional coordinate to index
idx = tf.concat([tf.expand_dims(tf.range(n), 1), k], axis=1)
# Make updates tensor
s = tf.shape(t2d_tile)
t2d_upd = tf.scatter_nd(idx, v, s)
# Make updates mask
upd_mask = tf.scatter_nd(idx, tf.ones_like(v, dtype=tf.bool), s)
# Make final tensor
t3d = tf.where(upd_mask, t2d_upd, t2d_tile)
# Test
with tf.Session() as sess:
print(sess.run(t3d))

输出：

[[[ 0  1 -2]
[ 3  4  5]]
[[ 0  1  2]
[-3  4  5]]]