jax.lax.scatter_mul#

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

散布-乘法运算符。

封装了 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 对象,描述 operandstart_indicesupdates 和输出的维度如何关联。

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

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

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

返回值:

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

返回类型:

数组