tf.tensor_scatter_nd_update

Разбросать updates по существующему тензору в соответствии с indices.

tf.tensor_scatter_nd_update(
    tensor, indices, updates, name=None
)

Эта операция создает новый тензор, применяя разреженные updates к входному tensor. Это похоже на присваивание по индексу.

# Not implemented: tensors cannot be updated inplace.
tensor[indices] = updates

Если на CPU обнаружен индекс за пределами допустимого диапазона, возвращается ошибка.

• Если обнаружен индекс за пределами допустимого диапазона, он игнорируется.
• Порядок применения обновлений не детерминирован, поэтому результат будет недетерминированным, если indices содержит дубликаты.

Эта операция очень похожа на tf.scatter_nd, за исключением того, что обновления записываются в существующий тензор (а не в тензор, заполненный нулями). Если память существующего тензора нельзя использовать повторно, создается его копия, в которую вносятся обновления.

В общем случае:

• indices — это целочисленный тензор, содержащий индексы для обновления в tensor.
• indices содержит как минимум две оси; последняя ось задает глубину векторов индексов.
• Каждому вектору индексов в indices соответствует элемент в updates.
• Если длина векторов индексов совпадает с рангом tensor, каждый вектор индексов указывает на скаляр в tensor, и каждое обновление является скаляром.
• Если длина векторов индексов меньше ранга tensor, каждый вектор индексов указывает на срез tensor, а форма обновлений должна совпадать с формой этого среза.

В итоге получаются следующие ограничения на формы:

assert tf.rank(indices) >= 2
index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]
assert index_depth <= tf.rank(tensor)
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == batch_shape + inner_shape

На практике обычно используется гораздо более простой вариант, чем эта общая форма. Его легче понять на простых примерах:

Скалярные обновления

Простейший вариант использования — вставка скалярных элементов в тензор по индексам. В этом случае index_depth должно быть равно рангу входного tensor, а срез каждого столбца indices является индексом оси входного tensor.

В этом простейшем случае ограничения на формы таковы:

num_updates, index_depth = indices.shape.as_list()
assert updates.shape == [num_updates]
assert index_depth == tf.rank(tensor)`

Например, чтобы вставить 4 элемента в тензор ранга 1, содержащий 8 элементов.

Операция scatter будет выглядеть так:

tensor = [0, 0, 0, 0, 0, 0, 0, 0]    # tf.rank(tensor) == 1
indices = [[1], [3], [4], [7]]       # num_updates == 4, index_depth == 1
updates = [9, 10, 11, 12]            # num_updates == 4
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor([ 0 9  0 10  11  0  0 12], shape=(8,), dtype=int32)

Длина (первая ось) updates должна совпадать с длиной indices: num_updates. Это количество вставляемых обновлений. Каждое скалярное обновление вставляется в tensor в указанную индексом позицию.

Для входного тензора tensor более высокого ранга скалярные обновления можно вставлять, используя index_depth, ранг которого совпадает с tf.rank(tensor):

tensor = [[1, 1], [1, 1], [1, 1]]    # tf.rank(tensor) == 2
indices = [[0, 1], [2, 0]]           # num_updates == 2, index_depth == 2
updates = [5, 10]                    # num_updates == 2
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor(
    [[ 1  5]
     [ 1  1]
     [10  1]], shape=(3, 2), dtype=int32)

Обновления срезов

Если входной tensor содержит более одной оси, scatter можно использовать для обновления целых срезов.

В этом случае удобно представить входной tensor как двухуровневый массив массивов. Форма этого двухуровневого массива разделяется на outer_shape и inner_shape.

indices индексирует внешний уровень входного тензора (outer_shape) и заменяет подмассив в указанной позиции соответствующим элементом списка updates. Форма каждого обновления — inner_shape.

При обновлении списка срезов ограничения на формы таковы:

num_updates, index_depth = indices.shape.as_list()
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == [num_updates, inner_shape]

Например, чтобы обновить строки (6, 3) tensor:

tensor = tf.zeros([6, 3], dtype=tf.int32)

Используйте глубину индекса 1.

indices = tf.constant([[2], [4]])     # num_updates == 2, index_depth == 1
num_updates, index_depth = indices.shape.as_list()

Значение outer_shape равно 6, внутренняя форма — 3:

outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]

Индексируются 2 строки, поэтому необходимо предоставить 2 updates. Форма каждого обновления должна совпадать с inner_shape.

# num_updates == 2, inner_shape==3
updates = tf.constant([[1, 2, 3],
                       [4, 5, 6]])

В итоге получаем:

tf.tensor_scatter_nd_update(tensor, indices, updates).numpy()
array([[0, 0, 0],
       [0, 0, 0],
       [1, 2, 3],
       [0, 0, 0],
       [4, 5, 6],
       [0, 0, 0]], dtype=int32)

Другие примеры обновления срезов

Тензор, представляющий пакет видеоклипов одинаковой длины, естественным образом имеет 5 осей: [batch_size, time, width, height, channels].

Например:

batch_size, time, width, height, channels = 13,11,7,5,3
video_batch = tf.zeros([batch_size, time, width, height, channels])

Чтобы заменить выбранные видеоклипы:

• Используйте index_depth глубины 1 (индексируя outer_shape: [batch_size])
• Предоставьте обновления, каждое из которых имеет форму, совпадающую с inner_shape: [time, width, height, channels].

Чтобы заменить первые два клипа единицами:

indices = [[0],[1]]
new_clips = tf.ones([2, time, width, height, channels])
tf.tensor_scatter_nd_update(video_batch, indices, new_clips)

Чтобы заменить выбранные кадры в видеоклипах:

• indices должен иметь index_depth глубины 2 для outer_shape: [batch_size, time].
• updates должен иметь форму списка изображений. Форма каждого обновления должна совпадать с inner_shape: [width, height, channels].

Чтобы заменить первый кадр первых трех видеоклипов:

indices = [[0, 0], [1, 0], [2, 0]] # num_updates=3, index_depth=2
new_images = tf.ones([
  # num_updates=3, inner_shape=(width, height, channels)
  3, width, height, channels])
tf.tensor_scatter_nd_update(video_batch, indices, new_images)

Свернутые индексы

В простых случаях удобно представлять indices и updates как списки, но это не является обязательным требованием. Вместо плоского num_updates значения indices и updates можно свернуть в batch_shape. Это batch_shape включает все оси indices, кроме самой внутренней оси index_depth.

index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]

Примечание: Единственное исключение: batch_shape нельзя []. Нельзя обновить один индекс, передав индексы с формой [index_depth].

updates должен иметь соответствующий batch_shape (оси перед inner_shape).

assert updates.shape == batch_shape + inner_shape

Примечание: Результат эквивалентен объединению осей batch_shape тензора indices и updates. Это обобщение позволяет избежать изменения формы, когда естественнее создавать «свернутые» индексы и обновления.

С учетом этого обобщения полные ограничения на формы выглядят так:

assert tf.rank(indices) >= 2
index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]
assert index_depth <= tf.rank(tensor)
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == batch_shape + inner_shape

Например, чтобы нарисовать X на матрице (5,5), начните со следующих индексов:

tensor = tf.zeros([5,5])
indices = tf.constant([
 [[0,0],
  [1,1],
  [2,2],
  [3,3],
  [4,4]],
 [[0,4],
  [1,3],
  [2,2],
  [3,1],
  [4,0]],
])
indices.shape.as_list()  # batch_shape == [2, 5], index_depth == 2
[2, 5, 2]

Здесь indices имеет форму не [num_updates, index_depth], а batch_shape+[index_depth].

Поскольку index_depth равно рангу tensor:

• outer_shape равно (5,5)
• inner_shape равно () — каждое обновление является скаляром
• updates.shape равно batch_shape + inner_shape == (5,2) + ()

updates = [
  [1,1,1,1,1],
  [1,1,1,1,1],
]

Объединив это, получаем:

tf.tensor_scatter_nd_update(tensor, indices, updates).numpy()
array([[1., 0., 0., 0., 1.],
       [0., 1., 0., 1., 0.],
       [0., 0., 1., 0., 0.],
       [0., 1., 0., 1., 0.],
       [1., 0., 0., 0., 1.]], dtype=float32)