jax.lax.scatter_min

目录

jax.lax.scatter_min#

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

Scatter-min 操作符。

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

scatter 的语义很复杂,其 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 是否已知为排序。如果是,可能会在某些后端提高性能。

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

  • mode (str | GatherScatterMode | None) – 如何处理越界索引:当设置为 ‘clip’ 时,索引会被限制在边界内,使得切片在边界内;当设置为 ‘fill’ 或 ‘drop’ 时,越界的更新会被丢弃。当设置为 ‘promise_in_bounds’ 时,越界索引的行为是实现定义的。

返回:

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

返回类型:

Array