Gather

Gathers elements of an input tensor into an output tensor.

Attributes

axis The axis to gather elements from, must obey \(0 \leq axis < rank(input)\).

mode The gather mode:

  • DEFAULT Similar to ONNX Gather. This is the default.

  • ELEMENT Similar to ONNX GatherElements.

  • ND Similar to ONNX GatherND.

num_elementwise_dims The dimension to start gathering from.

Inputs

input: tensor of type T1.

indices: tensor of type T2. Every element \(I_j\) in indices must obey \(0 \leq I_j < \text{input_dimensions}[axis]\)

Outputs

output: tensor of type T1.

Data Types

T1: int8, int32, int64, float16, float32, bfloat16

T2: int32, int64

  • if indices are int64 and the data is not int64, then the indices will be cast to int32.

Shape Information

input is a tensor with a shape of \([a_0,...,a_n]\).

indices is a tensor with a shape of \([i_0,...,i_m]\). When mode is set to ELEMENT, \(m=n\).

output is a tensor with a shape of:

  • when mode is set to DEFAULT: \([a_0,..., a_{axis-1}, i_{\text{num_elementwise_dims}},..., i_m, a_{axis+1},...,a_n]\), and its rank is n + m - 1 - num_elementwise_dims.

  • when mode is set to ELEMENTS: \([i_0,...,i_m]\).

  • when mode is set to ND: the rank of the tensor is \(n + m - i_m - 1 - axis\).

, where 0 < num_elementwise_dims <= 1.

Examples

Gather
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(1, 3))
layer = network.add_gather(in1, indices, axis=2)
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, -2.0, -1.0]]])
inputs[indices.name] = np.array([2, 1, 1])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[[-1.0, -2.0, -2.0]]], [[[2.0, 1.0, 1.0]]]])
Gather With 2D Indices
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 2))
layer = network.add_gather(in1, indices, axis=2)
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, -2.0, -1.0]]])
inputs[indices.name] = np.array([[2, 1], [3, 3]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[[-1.0, -2.0], [10.0, 10.0]]], [[[2.0, 1.0], [-2.0, -2.0]]]])
GatherElements
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ELEMENT)
layer.axis = 2
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, -2.0, -1.0]]])
inputs[indices.name] = np.array([[[0, 1, 2, 3, 4]], [[4, 2, 3, 3, 1]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[-1.0, 2.0, -2.0, -2.0, 1.0]]])
GatherElements With Axis=0
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ELEMENT)
layer.axis = 0
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, -2.0, -1.0]]])
inputs[indices.name] = np.array([[[0, 1, 1, 0, 0]], [[1, 1, 0, 0, 1]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[-3.0, 1.0, 2.0, 10.0, -25.0]], [[0.0, 1.0, -1.0, 10.0, -1.0]]])
GatherND
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 3))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ND)
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, -2.0, -1.0]]])
inputs[indices.name] = np.array([[0, 0, 2], [1, 0, 1]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[-1.0, 1.0]])

C++ API

For more information about the C++ IGatherLayer operator, refer to the C++ IGatherLayer documentation.

Python API

For more information about the Python IGatherLayer operator, refer to the Python IGatherLayer documentation.