jax.lax.scatter_max#

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

Scatter-max 操作符。

封装了 XLA 的 Scatter 操作符，其中 max 函数用于组合来自 operand 的更新和值。

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

参数:

  • operand (ArrayLike) – 一个应用 scatter 操作的数组。

  • 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 和散布更新总和的数组。

返回类型:

数组