jax.lax.scatter

jax.lax.scatter(operand, scatter_indices, updates, dimension_numbers, *, indices_are_sorted=False, unique_indices=False, mode=None)

散布更新运算符。

包装了 XLA 的 Scatter 运算符，其中更新会替换 operand 中的值。

如果对 operand 的同一个索引执行多次更新，则它们可以按任何顺序应用。

scatter() 是一个具有复杂语义的底层运算符，大多数 JAX 用户永远不需要直接调用它。相反，您应该优先使用 jax.numpy.ndarray.at() 来获得更熟悉的 NumPy 风格的索引语法。

参数:

• operand (ArrayLike) – 应该应用散布的数组

• scatter_indices (ArrayLike) – 一个数组，给出 operand 中每个 updates 中的更新应该应用到的索引。

• updates (ArrayLike) – 应该散布到 operand 上的更新。

• dimension_numbers (ScatterDimensionNumbers) – 一个 lax.ScatterDimensionNumbers 对象，描述 operand、start_indices、updates 和输出的维度之间是如何关联的。

• indices_are_sorted (bool) – 是否已知 scatter_indices 已排序。如果为 true，则可以提高某些后端的性能。

• unique_indices (bool) – 是否保证要更新的 operand 中的元素彼此不重叠。如果为 true，则可以提高某些后端的性能。JAX 不检查此承诺：如果当 unique_indices 为 True 时更新的元素重叠，则行为未定义。

• mode (str | GatherScatterMode | None | None) – 如何处理超出范围的索引：当设置为“clip”时，索引会被钳制，以便切片在边界内；当设置为“fill”或“drop”时，超出范围的更新将被丢弃。当设置为“promise_in_bounds”时，超出范围的索引的行为是实现定义的。

返回值:

一个包含 operand 和散布的更新之和的数组。

返回类型:

数组

示例

如上所述，您基本上永远不应该直接使用 scatter()，而是使用通过 jax.numpy.ndarray.at 的 NumPy 风格索引表达式执行散布式操作。

这是一个使用 jax.numpy.ndarray.at 更新数组中的条目的示例，它会降级为 XLA Scatter 操作

>>> x = jnp.zeros(5)
>>> indices = jnp.array([1, 2, 4])
>>> values = jnp.array([2.0, 3.0, 4.0])

>>> x.at[indices].set(values)
Array([0., 2., 3., 0., 4.], dtype=float32)

此语法还支持 scatter() 的几个可选参数，例如

>>> x.at[indices].set(values, indices_are_sorted=True, mode='promise_in_bounds')
Array([0., 2., 3., 0., 4.], dtype=float32)

相比之下，这是直接使用 scatter() 的等效函数调用，这不是典型用户需要做的

>>> lax.scatter(x, indices[:, None], values,
...             dimension_numbers=lax.ScatterDimensionNumbers(
...                 update_window_dims=(),
...                 inserted_window_dims=(0,),
...                 scatter_dims_to_operand_dims=(0,)),
...             indices_are_sorted=True,
...             mode=lax.GatherScatterMode.PROMISE_IN_BOUNDS)
Array([0., 2., 3., 0., 4.], dtype=float32)