jax.lax.scatter#

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

散佈更新運算子。

包裝 XLA 的 Scatter 運算子，其中更新會取代來自 operand 的值。

如果對 operand 的相同索引執行多個更新，則可能會以任何順序套用它們。

scatter() 是一個低階運算子，語義複雜，大多數 JAX 使用者永遠不需要直接呼叫它。相反地，您應該偏好使用 jax.numpy.ndarray.at() 以獲得更熟悉的 NumPy 風格索引語法。

參數:

  • operand (ArrayLike) – 應套用散佈的陣列

  • 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 和散佈更新總和的陣列。

回傳類型:

陣列

範例

如上所述，您基本上永遠不應該直接使用 scatter()，而是應該透過 jax.numpy.ndarray.at 使用 NumPy 風格的索引表示式來執行散佈風格的操作。

以下是使用 jax.numpy.ndarray.at 更新陣列中條目的範例，它會降低為 XLA Scatter 運算

>>> x = jnp.zeros(5)
>>> indices = jnp.array([1, 2, 4])
>>> values = jnp.array([2.0, 3.0, 4.0])

>>> x.at[indices].set(values)
Array([0., 2., 3., 0., 4.], dtype=float32)

此語法也支援 scatter() 的幾個可選引數，例如

>>> x.at[indices].set(values, indices_are_sorted=True, mode='promise_in_bounds')
Array([0., 2., 3., 0., 4.], dtype=float32)

相比之下，以下是直接使用 scatter() 的等效函式呼叫，這不是典型使用者應該需要做的事情

>>> lax.scatter(x, indices[:, None], values,
...             dimension_numbers=lax.ScatterDimensionNumbers(
...                 update_window_dims=(),
...                 inserted_window_dims=(0,),
...                 scatter_dims_to_operand_dims=(0,)),
...             indices_are_sorted=True,
...             mode=lax.GatherScatterMode.PROMISE_IN_BOUNDS)
Array([0., 2., 3., 0., 4.], dtype=float32)