jax.lax.scatter

内容

jax.lax.scatter#

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

散射更新操作符。

包装 XLA 的 Scatter 操作符，其中更新替换来自 operand 的值。

如果对操作数的相同索引执行多次更新，则可以按任意顺序应用它们。

散射的语义很复杂，并且其 API 可能会在将来发生变化。对于大多数用例，您应该优先考虑 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 和分散更新的总和的数组。

返回类型:

数组

内容