tf.compat.v1.gather_nd

Собирает срезы из params в тензор с формой, заданной indices.

Просмотреть псевдонимы

Совместимые псевдонимы для миграции

Подробнее см. в руководстве по миграции.

tf.compat.v1.manip.gather_nd

tf.compat.v1.gather_nd(
    params, indices, name=None, batch_dims=0
)

indices — это Tensor индексов в params. Векторы индексов расположены вдоль последней оси indices.

Это похоже на tf.gather, где indices задаёт срезы по первому измерению params. В tf.gather_nd indices задаёт срезы по первым N измерениям params, где N = indices.shape[-1].

Сбор скалярных значений

В простейшем случае векторы в indices индексируют полный ранг params:

tf.gather_nd(
    indices=[[0, 0],
             [1, 1]],
    params = [['a', 'b'],
              ['c', 'd']]).numpy()
array([b'a', b'd'], dtype=object)

В этом случае результат имеет на одну ось меньше, чем indices, а каждый вектор индексов заменяется скалярным значением, выбранным из params по индексу.

В этом случае соотношение форм имеет вид:

index_depth = indices.shape[-1]
assert index_depth == params.shape.rank
result_shape = indices.shape[:-1]

Если ранг indices равен K, полезно представлять indices как (K-1)-мерный тензор индексов в params.

Сбор срезов

Если векторы индексов не индексируют полный ранг params, то в каждой позиции результата содержится срез params. В этом примере собираются строки матрицы:

tf.gather_nd(
    indices = [[1],
               [0]],
    params = [['a', 'b', 'c'],
              ['d', 'e', 'f']]).numpy()
array([[b'd', b'e', b'f'],
       [b'a', b'b', b'c']], dtype=object)

Здесь indices содержит [2] векторов индексов длиной 1 каждый. Каждый вектор индексов указывает на строку матрицы params. Форма каждой строки — [3], поэтому форма результата — [2, 3].

В этом случае соотношение между формами имеет вид:

index_depth = indices.shape[-1]
outer_shape = indices.shape[:-1]
assert index_depth <= params.shape.rank
inner_shape = params.shape[index_depth:]
output_shape = outer_shape + inner_shape

В этом случае полезно представлять результат как тензоры, состоящие из тензоров. Форма внешнего тензора определяется ведущими измерениями indices. Форма внутренних тензоров совпадает с формой отдельного среза.

Пакеты

Кроме того, у params и indices может быть M совпадающих ведущих пакетных измерений. В этом случае для batch_dims необходимо задать значение M.

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

tf.gather_nd(
    indices = [[0, 1],
               [1, 0],
               [2, 4],
               [3, 2],
               [4, 1]],
    params=tf.zeros([5, 7, 3])).shape.as_list()
[5, 3]

Аргумент batch_dims позволяет опустить эти ведущие измерения расположения из индекса:

tf.gather_nd(
    batch_dims=1,
    indices = [[1],
               [0],
               [4],
               [2],
               [1]],
    params=tf.zeros([5, 7, 3])).shape.as_list()
[5, 3]

Это эквивалентно вызову отдельной операции gather_nd для каждого расположения в пакетных измерениях.

params=tf.zeros([5, 7, 3])
indices=tf.zeros([5, 1])
batch_dims = 1

index_depth = indices.shape[-1]
batch_shape = indices.shape[:batch_dims]
assert params.shape[:batch_dims] == batch_shape
outer_shape = indices.shape[batch_dims:-1]
assert index_depth <= params.shape.rank
inner_shape = params.shape[batch_dims + index_depth:]
output_shape = batch_shape + outer_shape + inner_shape
output_shape.as_list()
[5, 3]

Дополнительные примеры

Индексация 3-тензора:

tf.gather_nd(
    indices = [[1]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'a1', b'b1'],
        [b'c1', b'd1']]], dtype=object)

tf.gather_nd(
    indices = [[0, 1], [1, 0]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'c0', b'd0'],
       [b'a1', b'b1']], dtype=object)

tf.gather_nd(
    indices = [[0, 0, 1], [1, 0, 1]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([b'b0', b'b1'], dtype=object)

Примеры ниже относятся к случаю, когда дополнительные ведущие измерения есть только у индексов. Если у 'params' и 'indices' есть ведущие пакетные измерения, используйте параметр 'batch_dims', чтобы запустить gather_nd в пакетном режиме.

Пакетная индексация матрицы:

tf.gather_nd(
    indices = [[[0, 0]], [[0, 1]]],
    params = [['a', 'b'], ['c', 'd']]).numpy()
array([[b'a'],
       [b'b']], dtype=object)

Пакетная индексация срезов матрицы:

tf.gather_nd(
    indices = [[[1]], [[0]]],
    params = [['a', 'b'], ['c', 'd']]).numpy()
array([[[b'c', b'd']],
       [[b'a', b'b']]], dtype=object)

Пакетная индексация 3-тензора:

tf.gather_nd(
    indices = [[[1]], [[0]]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[[b'a1', b'b1'],
         [b'c1', b'd1']]],
       [[[b'a0', b'b0'],
         [b'c0', b'd0']]]], dtype=object)

tf.gather_nd(
    indices = [[[0, 1], [1, 0]], [[0, 0], [1, 1]]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'c0', b'd0'],
        [b'a1', b'b1']],
       [[b'a0', b'b0'],
        [b'c1', b'd1']]], dtype=object)

tf.gather_nd(
    indices = [[[0, 0, 1], [1, 0, 1]], [[0, 1, 1], [1, 1, 0]]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'b0', b'b1'],
       [b'd0', b'c1']], dtype=object)

Примеры с пакетными 'params' и 'indices':

tf.gather_nd(
    batch_dims = 1,
    indices = [[1],
               [0]],
    params = [[['a0', 'b0'],
               ['c0', 'd0']],
              [['a1', 'b1'],
               ['c1', 'd1']]]).numpy()
array([[b'c0', b'd0'],
       [b'a1', b'b1']], dtype=object)

tf.gather_nd(
    batch_dims = 1,
    indices = [[[1]], [[0]]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'c0', b'd0']],
       [[b'a1', b'b1']]], dtype=object)

tf.gather_nd(
    batch_dims = 1,
    indices = [[[1, 0]], [[0, 1]]],
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'c0'],
       [b'b1']], dtype=object)

См. также tf.gather.

Аргументы
`params`	`Tensor`. Тензор, из которого собираются значения.
`indices`	`Tensor`. Должен иметь один из следующих типов: `int32`, `int64`. Тензор индексов.
`name`	Имя операции (необязательно).
`batch_dims`	Целое число или скалярный объект 'Tensor'. Количество пакетных измерений.

Возвращает
`Tensor`. Имеет тот же тип, что и `params`.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/compat/v1/gather_nd