index_select

dragon.index_select(
  inputs,
  index,
  axis=0,
  **kwargs
)[source]

Select the elements according to the index along the given axis.

index could be a int64 tensor or a sequence with integers:

x = dragon.constant([[1, 2, 3], [4, 5, 6]])
print(dragon.index_select(x, [0, 1]))
print(dragon.index_select(x, dragon.constant([0, 1], 'int64')))

More than one axis could be specified for index:

# The number of ``axis`` should less than ``rank(index)``
# And these axes should be continuous
print(dragon.index_select(x, [0, 1], axis=[0, 1]))
Parameters:
  • inputs (dragon.Tensor) – The input tensor.
  • index (Union[Sequence[int], dragon.Tensor]) – The index to select elements.
  • axis (Union[int, Sequence[int]], optional, default=0) – The axis where the index aligned.
Returns:

dragon.Tensor – The output tensor.