Расширенное индексирование Tensorflow: назначьте меньший тензор в больший тензор в позицию, основанную на двух индексных тензорах

Question

Я хотел бы создать тензор нулей с формой (3,4,2,2) и вставить в него тензор (3,4) в позиции, заданной двумя (3,1) тензорами.

Пример кода: Эквивалентная операция NumPy для массивов будет выглядеть следующим образом:

# Existing arrays of required shapes
bbox = np.arange(3*4).reshape(3,4)
x = np.array([0,0,1])
y = np.array([1,1,1])

# Create zeros array and assign into it
output = np.zeros((3,4,2,2))
output[np.arange(3),:,x,y] = bbox

Как мне сделать что-то похожее с Tensorflow?

ПРИМЕЧАНИЕ. На самом деле я хочу работать с тензоромразмером (32,125,32,32).Выше приведен простой код для воспроизведения

Answer 1

Вот как это можно сделать с помощью tf.scatter_nd:

import tensorflow as tf
import numpy as np

bbox = np.arange(3 * 4).reshape(3, 4)
x = np.array([0, 0, 1])
y = np.array([1, 1, 1])
x_size = 2
y_size = 2

# TensorFlow calculation
with tf.Graph().as_default(), tf.Session() as sess:
    bbox_t = tf.convert_to_tensor(bbox)
    x_t = tf.convert_to_tensor(x)
    y_t = tf.convert_to_tensor(y)
    shape = tf.shape(bbox_t)
    rows, cols = shape[0], shape[1]
    ii, jj = tf.meshgrid(tf.range(rows), tf.range(cols), indexing='ij')
    xx = tf.tile(tf.expand_dims(x_t, 1), (1, cols))
    yy = tf.tile(tf.expand_dims(y_t, 1), (1, cols))
    idx = tf.stack([ii, jj, xx, yy], axis=-1)
    output = tf.scatter_nd(idx, bbox_t, [rows, cols, x_size, y_size])
    output_tf = sess.run(output)

# Test with NumPy calculation
output_np = np.zeros((3, 4, 2, 2))
output_np[np.arange(3), :, x, y] = bbox
print(np.all(output_tf == output_np))
# True