tf.raw_ops.GatherV2

Извлечение срезов из params по оси axis в соответствии с indices.

tf.raw_ops.GatherV2(
    params, indices, axis, batch_dims=0, name=None
)

indices должно быть целочисленным тензором любой размерности (обычно 0-й или 1-й). Создаёт выходной тензор с формой params.shape[:axis] + indices.shape[batch_dims:] + params.shape[axis + 1:], где:

# Scalar indices (output is rank(params) - 1).
output[a_0, ..., a_n, b_0, ..., b_n] =
  params[a_0, ..., a_n, indices, b_0, ..., b_n]

# Vector indices (output is rank(params)).
output[a_0, ..., a_n, i, b_0, ..., b_n] =
  params[a_0, ..., a_n, indices[i], b_0, ..., b_n]

# Higher rank indices (output is rank(params) + rank(indices) - 1).
output[a_0, ..., a_n, i, ..., j, b_0, ... b_n] =
  params[a_0, ..., a_n, indices[i, ..., j], b_0, ..., b_n]

Обратите внимание, что на процессоре (CPU), если найден индекс вне границ, возвращается ошибка. На графическом процессоре (GPU), если найден индекс вне границ, в соответствующее значение выходного тензора записывается 0.

См. также tf.batch_gather и tf.gather_nd.