Scatter

Creates an output tensor by copying values from an input tensor and then updating values by the given indices and updates tensors.

Attributes

mode The scatter mode to use. Can be one of the following:

• ELEMENT indices and tensors defines an update action on elements of the input.

• ND indices and tensors defines an update action on slices of the input.

axis The axis to scatter on. Default to 0.

Inputs

input: tensor of type T1.

indices: tensor of type T2.

updates: tensor of type T1.

Outputs

output: tensor of type T1.

Data Types

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

T2: int32, int64

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

updates is a tensor with a shape of \([j_0,...,j_k]\). When mode is set to ELEMENT, \(k=n\). When mode is set to ND, \(k = n + m - i_m - 1 - 1\).

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

Volume Limits

input, indices, and updates can have up to \(2^{31}-1\) elements.

Examples

Scatter

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1, 5))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.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([[[1, 2, 3, 0, 4]], [[4, 1, 2, 3, 0]]])

inputs[updates.name] = np.array([[[-1.0, 2.4, 3.2, 10.8, 8.9]], [[0, -11.2, 34.2, 223.9, -100]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[10.8, -1.0, 2.4, 3.2, 8.9]], [[-100, -11.2, 34.2, 223.9, 0]]])

Scatter On 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))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1, 5))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.ELEMENT)
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([[[1, 0, 0, 0, 1]], [[0, 1, 1, 1, 0]]])

inputs[updates.name] = np.array([[[-1.0, 2.4, 3.2, 10.8, 8.9]], [[0, -11.2, 34.2, 223.9, -100]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[0, 2.4, 3.2, 10.8, -100]], [[-1.0, -11.2, 34.2, 223.9, 8.9]]])

ScatterND

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 3))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.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, 1]]], [[[1, 0, 3]]]])

inputs[updates.name] = np.array([[-1.0, 0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[-3.0, -1.0, -1.0, 10.0, -25.0]], [[0.0, 1.0, 2.0, 0.0, -1.0]]])

C++ API

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

Python API

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