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

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.