jax.lax.scatter#
- jax.lax.scatter(operand, scatter_indices, updates, dimension_numbers, *, indices_are_sorted=False, unique_indices=False, mode=None)[源代码]#
Scatter-update 运算符。
包装了 XLA 的 Scatter 运算符,其中更新替换了 operand 中的值。
如果对 operand 的相同索引执行多次更新,则可以按任何顺序应用它们。
scatter()
是一个低级运算符,语义复杂,大多数 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 更新之和的数组。
- 返回类型:
示例
如上所述,您基本上永远不应该直接使用
scatter()
,而是应该使用 NumPy 风格的索引表达式,通过jax.numpy.ndarray.at
执行 scatter 风格的操作。这是一个使用
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)