jax.ref.new_ref

jax.ref.new_ref(init_val, *, memory_space=None)

Create a mutable array reference with initial value init_val.

For more discussion, see the Ref guide.

Parameters:

• init_val (Any) – A jax.Array representing the initial state of the buffer.

• memory_space (Any) – An optional memory space attribute for the Ref.

Returns:

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

Return type:

Ref