add HijaxVariable

Author: cgarciae

docs_nnx/api_reference/flax.nnx/graph.rst

@@ -31,7 +31,8 @@ graph
 .. autofunction:: find_duplicates
 .. autofunction:: pure
 .. autofunction:: as_ref
-.. autofunction:: as_hijax
-.. autofunction:: as_lojax
+.. autofunction:: as_mutable
+.. autofunction:: as_immutable
+.. autofunction:: as_pytree
 .. autofunction:: flatten
 .. autofunction:: unflatten

docs_nnx/guides/hijax.ipynb

@@ -0,0 +1,58 @@
+Hijax (experimental)
+====================
+
+
+
+----
+
+Basic usage
+^^^^^^^^^^^^
+
+.. testsetup::
+
+  import jax
+  import jax.numpy as jnp
+
+  current_mode = nnx.current_variable_mode()
+
+.. testcode::
+
+  from flax import nnx
+  import optax
+
+  nnx.variable_mode('mutable') # | 'ref'
+
+  class Model(nnx.Module):
+    def __init__(self, din, dmid, dout, rngs: nnx.Rngs):
+      self.linear = nnx.Linear(din, dmid, rngs=rngs)
+      self.bn = nnx.BatchNorm(dmid, rngs=rngs)
+      self.dropout = nnx.Dropout(0.2)
+      self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)
+
+    def __call__(self, x, rngs):
+      x = nnx.relu(self.dropout(self.bn(self.linear(x)), rngs=rngs))
+      return self.linear_out(x)
+
+  model = Model(2, 64, 3, rngs=nnx.Rngs(0))  # eager initialization
+  optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
+
+  @jax.jit
+  def train_step(model, optimizer, rngs, x, y):
+    graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)
+    def loss_fn(params):
+      model = nnx.merge(graphdef, params, nondiff)
+      return ((model(x, rngs) - y) ** 2).mean()
+    loss, grads = jax.value_and_grad(loss_fn)(nnx.as_immutable(params))
+    optimizer.update(model, grads)  # in-place updates
+    return loss
+
+  nnx.variable_mode(current_mode)  # clean up for CI tests
+
+
+----
+
+.. toctree::
+   :hidden:
+   :maxdepth: 2
+
+   variable