jax.ref module

jax.ref has the API for working with ArrayRef.

API

AbstractRef(inner_aval[, memory_space, kind])	Abstract mutable array reference.
Ref(aval, refs)	Mutable array reference.
freeze(ref)	Invalidate a given reference and return its final value.
get(ref[, idx])	Read a value from an Ref.
new_ref(init_val, *[, memory_space])	Create a mutable array reference with initial value init_val.
set(ref, idx, value)	Set a value in an Ref in-place.
swap(ref, idx, value[, _function_name])	Update an array value inplace while returning the previous value.
addupdate(ref, idx, x)	Add to an element in an Ref in-place.