jax.lax.scatter_max

内容

jax.lax.scatter_max#

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

散射-最大运算符。

包装 XLA 的 Scatter 运算符，其中使用 `max` 函数来组合更新和来自 `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 已排序。 如果为真，则可能提高某些后端的性能。

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

  • mode (str | GatherScatterMode | None | None) – 处理超出范围索引的方式：当设置为“clip”时，索引将被截断，以使切片在范围内；当设置为“fill”或“drop”时，超出范围的更新将被丢弃。 当设置为“promise_in_bounds”时，超出范围索引的行为是实现定义的。

返回值:

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

返回类型:

数组

内容