jax.Array.at#
- 抽象 属性 Array.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)