使用tf.where（）通过2d条件选择3d张量并用键和值替换2d索引中的元素

Question

There are 2 questions in the title. I am confused by both questions because tensorflow is such a static programming language (I really want to go back to either pytorch or chainer).

I give 2 examples. please answer me in tensorflow codes or providing the relevant function links.

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

(I don't know how to broadcast cond to 3d tensor)

2) change element in 2d tensor

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

Thank you so much in advance. XD

Answer 1

This are two quite different questions, and they should probably have been posted as such, but anyway.

1)

Yes, you need to manually broadcast all the inputs to [ tf.where ]( https://www.tensorflow.org/api_docs/python/tf/where] if they are different. For what is worth, there is an (old) open issue about it , but so far implicit broadcasting it has not been implemented. You can use tf.stack like you suggest, although tf.tile would probably be more obvious (and may save memory, although I'm not sure how it is implemented really):

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

Or simply with tf.broadcast_to :

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

2)

This is one way to do that:

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

Output:

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

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