jax.Array.at#
- abstract property 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)