Переменная тензор потока назначить трансляцию

Question

Есть ли какой-нибудь метод в tenorflow, позволяющий добиться широковещательного присвоения матрице (tf.Variable) что-то вроде следующего кода....

a    = tf.Variable(np.zeros([10,10,10,10], np.int32))

# creating a mask and trying to assign the 2nd, 3rd dimension of a
mask = tf.ones([10,10])

# 1) which is work in this case, but only assign one block
op = a[0,:,:,0].assign(mask)

# 2) attempting to broadcasting while not work, size mismatch
op = a[0].assign(mask)

для меня текущее решение может повторять все остальные измерения, но может пострадать от вложенного цикла, как в 1), или должен быть более разумный способ сделать это, спасибо!

Answer 1

Не общее решение (множество жестко закодированных тензорных форм), но, надеюсь, это даст вам суть для вашего примера:

a = tf.Variable(np.zeros([10,10,10,10], np.int32))
mask = tf.ones([10,10],dtype=tf.int32)
mask_reshaped = tf.reshape(mask,[1,10,10,1]) # make the number of dims match
mask_broadcast = tf.tile(mask_reshaped, [10, 1, 1, 10]) # do the actual broadcast
op = a.assign(mask_broadcast)