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: bool, int4, int8, int32, int64, float8, 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.