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.