jax.lax.scatter_mul#

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

Scatter-multiply 运算符。

包装 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、start_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 和分散的更新的乘积。

返回类型:

Array