jax.lax.scatter_apply#

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

Scatter-apply 運算子。

包裝 XLA 的 Scatter 運算子，其中來自 operand 的值會被替換為 func(operand)，重複的索引會導致多次應用 func。

scatter 的語義相當複雜，其 API 未來可能會變更。對於大多數使用案例，您應該偏好 JAX 陣列上的 jax.numpy.ndarray.at 屬性，它使用熟悉的 NumPy 索引語法。

請注意，在目前的實作中，scatter_apply 與自動微分不相容。

參數:

  • operand (Array) – 應套用 scatter 的陣列

  • scatter_indices (Array) – 一個陣列，提供 operand 中的索引，updates 中的每個更新都應套用至該索引。

  • func (Callable[[Array], Array]) – 將在每個索引套用的一元函式。

  • dimension_numbers (ScatterDimensionNumbers) – 一個 lax.ScatterDimensionNumbers 物件，描述 operand、start_indices、updates 和輸出的維度如何相關。

  • update_shape (Shape) – 給定索引處更新的形狀。

  • 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’ 時，超出範圍索引的行為是實作定義的。

傳回:

一個陣列，包含在給定索引處將 func 套用至 operand 的結果。

傳回類型:

Array