gather

dragon.gather(
  inputs,
  axis=0,
  end_axis=None,
  **kwargs
)[source]

Gather elements along the given axis using index.

Index should be a int64 tensor:

input = dragon.constant([[1, 2, 3], [4, 5, 6]])
index = dragon.constant([1])
print(dragon.gather([input, index]))  # [[4, 5, 6]]

More than one axis could be specified to gather:

# Along the continuous axes: [axis, end_axis]
print(dragon.gather([input, index], axis=0, end_axis=1))
Parameters:
  • inputs (Sequence[dragon.Tensor]) – The input and index tensor.
  • axis (int, optional, default=0) – The first axis to gather.
  • end_axis (int, optional) – The last axis to gather.
Returns:

dragon.Tensor – The output tensor.