What does this PR do?

• Variable(..., is_hijax=True) can now be used to HijaxVariable instances, which is type registered in the new Hijax system. HijaxVariable is a Mutable Hijax type meaning its instances will automatically propagate state through JAX transforms in a similar manner to jax.Ref.

• nnx.using_refs and nnx.use_refs were dropped in favor of nnx.using_hijax and nnx.use_hijax.

• Variable(..., is_mutable=False) can now be used to disallow any form of updates on the Variable instances. When is_mutable=False then neither is_hijax=True or has_ref=True are allowed.

• Added APIs to create an updated copy of a Variable or structure of Variables with a new value for is_hijax, has_ref, or is_mutable:

  • as_immutable_vars: create copies of Variables with is_mutable set to False. is_hijax or had_ref will be set to False as well if they had been active.
  • as_mutable_vars: create copies of Variables with is_mutable set to True. is_hijax or had_ref will be set back to True if they had been if as_immutable_vars as previously set them to False.
  • as_hijax_vars: create copies of Variables with is_hijax set to True.
  • as_pytree_vars: create copies of Variables with is_hijax set to False.
  • as_ref_vars: create copies of Variables with has_ref set to True.
  • as_array_vars: create copies of Variables with has_ref set to False.

  Note: there is currently some coupling between is_hijax, has_ref, or is_mutable which will be removed in the near future as HijaxVariable improves.