jax.ref.array_ref

Contents

jax.ref.array_ref#

jax.ref.array_ref(init_val, *, memory_space=None)[source]#

Create a mutable array reference with initial value init_val.

For more discussion, see the ArrayRef guide.

Parameters:

init_val – A jax.Array representing the initial state of the buffer.
memory_space (Any) – An optional memory space attribute for the ArrayRef.

Returns:

A jax.ref.ArrayRef containing a reference to a mutable buffer.

Return type: