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.
Outputs¶
output: tensor of type T1
.
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.