Skip to main content

Mojo function

scatter_set_constant

scatter_set_constant[data_type: DType, index_type: DType, //, target: StringSlice[StaticConstantOrigin], single_thread_blocking_override: Bool = False](data: LayoutTensor[data_type, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], indices: LayoutTensor[index_type, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], fill_value: SIMD[data_type, 1], ctx: DeviceContextPtr)

Scatter the fill_value into the data at the specified indices.

Example: Suppose we have a 3x3 matrix data initialized to zeros:

data = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]

And indices is a 2D tensor with shape [2, 2]:

indices = [[0, 1], [2, 0]]

If fill_value is 5, after calling scatter_set_constant, data will be:

data = [[0, 5, 0], [0, 0, 0], [5, 0, 0]]

Arguments: data: The data to scatter the updates into. indices: The indices to scatter the updates into. fill_value: The value to fill the data with. ctx: The device context.

Was this page helpful?