gather

dragon.vm.tensorflow.gather(
  params,
  indices,
  axis=0,
  name=None
)[source]

Gather the elements along the given axis using index.

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

x = tf.constant([[1, 2, 3], [4, 5, 6]])
print(tf.gather(x, [0, 1]))
print(tf.gather(x, tf.constant([0, 1], tf.int64)))
Parameters:
  • params (dragon.Tensor) – The tensor to provide elements.
  • indices (Union[Sequence[int], dragon.Tensor]) – The indices to gather.
  • axis (Union[int, Sequence[int]], optional, default=0) – The axis to gather.
  • name (str, optional) – The operation name.
Returns:

dragon.Tensor – The output tensor.