jax.ref 模块

jax.ref 提供了用于操作 ArrayRef 的 API。

API

AbstractRef(inner_aval[, memory_space, kind])	抽象可变数组引用。
Ref(aval, refs)	可变数组引用。
freeze(ref)	使给定的引用失效并返回其最终值。
get(ref[, idx])	从 Ref 读取一个值。
new_ref(init_val, *[, memory_space, kind])	创建一个初始值为 init_val 的可变数组引用。
set(ref, idx, value)	在原位（in-place）设置 Ref 中的值。
swap(ref, idx, value[, _function_name])	在原位更新数组值并返回之前的值。
addupdate(ref, idx, x)	在原位对 Ref 中的元素进行加法更新。