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) – 應 scatter 到 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 和 scatter 更新總和的陣列。

返回類型:

Array