jax.lax.scatter_apply

jax.lax.scatter_apply#

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

散点应用运算符。

封装了 XLA 的 Scatter 运算符，其中来自 operand 的值被替换为 func(operand)，重复索引会导致 func 的多次应用。

Scatter 的语义很复杂，其 API 未来可能会发生变化。对于大多数用例，您应该优先考虑 JAX 数组上的 jax.numpy.ndarray.at 属性，它使用熟悉的 NumPy 索引语法。

请注意，在当前实现中，scatter_apply 与自动微分不兼容。

参数:

  • operand (Array) – 应用 scatter 的数组。

  • scatter_indices (Array) – 一个数组，提供 operand 中每个更新应该应用到的索引。

  • func (Callable[[Array], Array]) – 将在每个索引上应用的一元函数。

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

  • update_shape (Shape) – 在给定索引处的更新形状。

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

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

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

返回值:

一个数组，包含在给定索引处将 func 应用于 operand 的结果。

返回类型:

Array

内容