jax.numpy.ndarray.at#

abstract property ndarray.at[源代码]#

用于索引更新功能的辅助属性。

at 属性提供了与就地数组修改功能上等效的纯函数式方法。

具体来说：

替代语法

等效的就地表达式

x = x.at[idx].set(y)

x[idx] = y

x = x.at[idx].add(y)

x[idx] += y

x = x.at[idx].subtract(y)

x[idx] -= y

x = x.at[idx].multiply(y)

x[idx] *= y

x = x.at[idx].divide(y)

x[idx] /= y

x = x.at[idx].power(y)

x[idx] **= y

x = x.at[idx].min(y)

x[idx] = minimum(x[idx], y)

x = x.at[idx].max(y)

x[idx] = maximum(x[idx], y)

x = x.at[idx].apply(ufunc)

ufunc.at(x, idx)

x = x.at[idx].get()

x = x[idx]

任何 x.at 表达式都不会修改原始的 x；相反，它们会返回 x 的修改后的副本。然而，在 jit() 编译的函数内部，像 x = x.at[idx].set(y) 这样的表达式保证会就地应用。

与 NumPy 的就地操作（如 x[idx] += y）不同，如果多个索引指向同一位置，则所有更新都将被应用（NumPy 只会应用最后一个更新，而不是应用所有更新）。冲突更新的应用顺序是实现定义的，并且可能是非确定性的（例如，由于某些硬件平台上的并发）。

默认情况下，JAX 假设所有索引都在界内。可以通过 mode 参数指定替代的越界索引语义（请参见下文）。

参数:

  • mode (str) –

    指定越界索引模式。选项包括：

    • "promise_in_bounds"：(默认) 用户保证索引在界内。不会执行额外的检查。实际上，这意味着 get() 中的越界索引将被裁剪，set()、add() 等中的越界索引将被丢弃。

    • "clip"：将越界索引钳制到有效范围内。

    • "drop"：忽略越界索引。

    • "fill"："drop" 的别名。对于 get()，可选的 fill_value 参数指定将返回的值。

      有关更多详细信息，请参阅 jax.lax.GatherScatterMode。

  • indices_are_sorted (bool) – 如果为 True，则实现将假设传递给 at[] 的索引按升序排序，这可以在某些后端带来更高效的执行。

  • unique_indices (bool) – 如果为 True，则实现将假设传递给 at[] 的索引是唯一的，这可以提高在某些后端上的执行效率。

  • fill_value (Any) – 仅适用于 get() 方法：当 mode 为 'fill' 时，为越界切片返回的填充值。否则将被忽略。对于非精确类型，默认为 NaN，对于有符号类型，默认为最大的负值，对于无符号类型，默认为最大的正值，对于布尔值，默认为 True。

示例

>>> x = jnp.arange(5.0)
>>> x
Array([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[2].add(10)
Array([ 0.,  1., 12.,  3.,  4.], dtype=float32)
>>> x.at[10].add(10)  # out-of-bounds indices are ignored
Array([0., 1., 2., 3., 4.], dtype=float32)
>>> x.at[20].add(10, mode='clip')
Array([ 0.,  1.,  2.,  3., 14.], dtype=float32)
>>> x.at[2].get()
Array(2., dtype=float32)
>>> x.at[20].get()  # out-of-bounds indices clipped
Array(4., dtype=float32)
>>> x.at[20].get(mode='fill')  # out-of-bounds indices filled with NaN
Array(nan, dtype=float32)
>>> x.at[20].get(mode='fill', fill_value=-1)  # custom fill value
Array(-1., dtype=float32)