Warning

This page was created from a pull request.

jax.lax.scatter

jax.lax.scatter(operand, scatter_indices, updates, dimension_numbers, *, indices_are_sorted=False, unique_indices=False)[source]

Scatter-update operator.

Wraps XLA’s Scatter operator, where updates replace values from operand.

If multiple updates are performed to the same index of operand, they may be applied in any order.

The semantics of scatter are complicated and its API is subject to change.

Parameters

  • operand (Any) – an array to which the scatter should be applied

  • scatter_indices (Any) – an array that gives the indices in operand to which each update in updates should be applied.

  • updates (Any) – the updates that should be scattered onto operand.

  • dimension_numbers (ScatterDimensionNumbers) – a lax.ScatterDimensionNumbers object that describes how dimensions of operand, start_indices, updates and the output relate.

  • indices_are_sorted (bool) – whether scatter_indices is known to be sorted. If true, may improve performance on some backends.

  • unique_indices (bool) – whether the indices to be updated in operand are guaranteed to not overlap with each other. If true, may improve performance on some backends.

Return type

Any

Returns

An array containing the sum of operand and the scattered updates.