jax.lax.scatter_add#

jax.lax.scatter_add(operand, scatter_indices, updates, dimension_numbers, *, indices_are_sorted=False, unique_indices=False, mode=None)[源代码]#

Scatter-add 运算符。

封装了 XLA 的 Scatter 运算符，其中加法用于组合来自 operand 的更新和值。

scatter 的语义很复杂，并且其 API 将来可能会发生变化。对于大多数用例，您应该首选 JAX 数组上的 jax.numpy.ndarray.at 属性，它使用熟悉的 NumPy 索引语法。

参数:

  • operand (ArrayLike) – 要应用分散的数组

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

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

  • dimension_numbers (ScatterDimensionNumbers) – 一个 lax.ScatterDimensionNumbers 对象，描述了 operand、scatter_indices、updates 和输出的维度之间的关系。

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

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

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

返回:

一个数组，包含 operand 和分散的更新的总和。

返回类型:

数组

示例

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

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

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

>>> x.at[indices].add(values)
Array([1., 3., 4., 1., 5.], dtype=float32)

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

>>> x.at[indices].add(values, indices_are_sorted=True,
...                   mode='promise_in_bounds')
Array([1., 3., 4., 1., 5.], dtype=float32)

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

>>> lax.scatter_add(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([1., 3., 4., 1., 5.], dtype=float32)