gather

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

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

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)))

More than one axis could be specified for indices:

# The number of ``axis`` should less than rank(indices)
# And these axes should be continuous
print(tf.gather(x, [0, 1], axis=[0, 1]))
Parameters:
  • params (dragon.Tensor) – The tensor to provide elements.
  • indices (Union[Sequence[int], dragon.Tensor]) – The indices of select.
  • axis (Union[int, Sequence[int]], optional, default=0) – The axis where the indices aligned.
  • name (str, optional) – A optional name for the operation.
Returns:

dragon.Tensor – The output tensor.