jax.lax.scatter_add#

jax.lax.scatter_add(operand, scatter_indices, updates, dimension_numbers, *, indices_are_sorted=False, unique_indices=False, mode=None)[原始碼]#

Scatter-add 運算子。

包裝了 XLA 的 Scatter 運算子，其中加法用於組合來自 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、scatter_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