Rename NNX view functions to have common naming convention.

Author: samanklesaria

docs_nnx/guides/view.ipynb

@@ -6,7 +6,7 @@
    "metadata": {},
    "source": [
     "# Model Views\n",
-    "This guide covers how to use the `nnx.view` function. This function is useful for handling state in layers like `Dropout` and `BatchNorm`, which behave differently in training and evaluation. Similar to `.view` for numpy arrays, `nnx.view` allows you to set modes of the model while still sharing the same data. For a quick intro to how this function works, refer to the following example:"
+    "This guide covers how to use NNX \"Views\", which are useful for handling state in layers like `Dropout` and `BatchNorm` which behave differently in training and evaluation. Similar to `.view` for numpy arrays, NNX views allow you to modify static attributes of the model while still sharing the same data. For a quick intro, consider the following example showcasing `nnx.with_modules`, an NNX View that overwrites module attributes."
    ]
   },
   {
@@ -25,8 +25,8 @@
     ")\n",
     "\n",
     "# set train and eval modes\n",
-    "train_model = nnx.view(model, deterministic=False, use_running_average=False)\n",
-    "eval_model = nnx.view(model, deterministic=True, use_running_average=True)\n",
+    "train_model = nnx.with_modules(model, deterministic=False, use_running_average=False)\n",
+    "eval_model = nnx.with_modules(model, deterministic=True, use_running_average=True)\n",
     "\n",
     "# Can see deterministic is different between train_model and eval_model\n",
     "assert train_model.layers[2].deterministic == False\n",
@@ -35,7 +35,7 @@
     "# Weights are shared between the models\n",
     "assert train_model.layers[0].kernel is eval_model.layers[0].kernel\n",
     "\n",
-    "# Print information about kwargs for nnx.view with nnx.view_info\n",
+    "# Print information about kwargs for nnx.with_modules with nnx.view_info\n",
     "print(nnx.view_info(model))"
    ]
   },
@@ -125,8 +125,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "train_model = nnx.view(model, deterministic=False)\n",
-    "eval_model = nnx.view(model, deterministic=True)\n",
+    "train_model = nnx.with_modules(model, deterministic=False)\n",
+    "eval_model = nnx.with_modules(model, deterministic=True)\n",
     "\n",
     "# weights are references to the same data\n",
     "assert train_model.lin1.kernel is eval_model.lin1.kernel\n",
@@ -196,8 +196,8 @@
    "source": [
     "model = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=rngs)\n",
     "optimizer = nnx.Optimizer(model, optax.adam(lr), wrt=nnx.Param)\n",
-    "train_model = nnx.view(model, deterministic=False)  # training view\n",
-    "eval_model = nnx.view(model, deterministic=True)  # eval view\n",
+    "train_model = nnx.with_modules(model, deterministic=False)  # training view\n",
+    "eval_model = nnx.with_modules(model, deterministic=True)  # eval view\n",
     "\n",
     "eval_results = []\n",
     "for epoch in range(total_epochs):\n",
@@ -293,7 +293,7 @@
     "\n",
     "\n",
     "model = PrintLayer()\n",
-    "model_print = nnx.view(model, msg='Hello, World!')\n",
+    "model_print = nnx.with_modules(model, msg='Hello, World!')\n",
     "\n",
     "model() # nothing printed\n",
     "model_print() # prints \"Hello, World!\""

docs_nnx/guides/view.md

@@ -9,7 +9,7 @@ jupytext:
 ---
 
 # Model Views
-This guide covers how to use the `nnx.view` function. This function is useful for handling state in layers like `Dropout` and `BatchNorm`, which behave differently in training and evaluation. Similar to `.view` for numpy arrays, `nnx.view` allows you to set modes of the model while still sharing the same data. For a quick intro to how this function works, refer to the following example:
+This guide covers how to use NNX "Views", which are useful for handling state in layers like `Dropout` and `BatchNorm` which behave differently in training and evaluation. Similar to `.view` for numpy arrays, NNX views allow you to modify static attributes of the model while still sharing the same data. For a quick intro, consider the following example showcasing `nnx.with_modules`, an NNX View that overwrites module attributes.
 
 ```{code-cell}
 from flax import nnx
@@ -21,8 +21,8 @@ model = nnx.Sequential(
 )
 
 # set train and eval modes
-train_model = nnx.view(model, deterministic=False, use_running_average=False)
-eval_model = nnx.view(model, deterministic=True, use_running_average=True)
+train_model = nnx.with_modules(model, deterministic=False, use_running_average=False)
+eval_model = nnx.with_modules(model, deterministic=True, use_running_average=True)
 
 # Can see deterministic is different between train_model and eval_model
 assert train_model.layers[2].deterministic == False
@@ -31,7 +31,7 @@ assert eval_model.layers[2].deterministic == True
 # Weights are shared between the models
 assert train_model.layers[0].kernel is eval_model.layers[0].kernel
 
-# Print information about kwargs for nnx.view with nnx.view_info
+# Print information about kwargs for nnx.with_modules with nnx.view_info
 print(nnx.view_info(model))
 ```
 
@@ -85,8 +85,8 @@ From the model display, we can see that `Dropout` has `deterministic == False`,
 This is where `nnx.view` comes in. This function updates the modes for each submodule of a neural network based on the kwargs passed into the function. The underlying model weights are then shared between different views. We set up a training and evaluation version of the model below.
 
 ```{code-cell}
-train_model = nnx.view(model, deterministic=False)
-eval_model = nnx.view(model, deterministic=True)
+train_model = nnx.with_modules(model, deterministic=False)
+eval_model = nnx.with_modules(model, deterministic=True)
 
 # weights are references to the same data
 assert train_model.lin1.kernel is eval_model.lin1.kernel
@@ -128,8 +128,8 @@ Now we create `train_model` and `eval_model` views up front. During the training
 ```{code-cell}
 model = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=rngs)
 optimizer = nnx.Optimizer(model, optax.adam(lr), wrt=nnx.Param)
-train_model = nnx.view(model, deterministic=False)  # training view
-eval_model = nnx.view(model, deterministic=True)  # eval view
+train_model = nnx.with_modules(model, deterministic=False)  # training view
+eval_model = nnx.with_modules(model, deterministic=True)  # eval view
 
 eval_results = []
 for epoch in range(total_epochs):
@@ -201,7 +201,7 @@ class PrintLayer(nnx.Module):
 
 
 model = PrintLayer()
-model_print = nnx.view(model, msg='Hello, World!')
+model_print = nnx.with_modules(model, msg='Hello, World!')
 
 model() # nothing printed
 model_print() # prints "Hello, World!"

docs_nnx/hijax/hijax.ipynb

@@ -49,7 +49,7 @@
     "@jax.jit\n",
     "def train_step(x, y):\n",
     "  loss_fn = lambda m: jnp.mean((m(x) - y) ** 2)\n",
-    "  loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(model, mutable=False))  # tmp fix for jax.grad\n",
+    "  loss, grads = jax.value_and_grad(loss_fn)(nnx.with_vars(model, mutable=False))  # tmp fix for jax.grad\n",
     "  optimizer.update(model, grads)\n",
     "  return loss\n",
     "\n",
@@ -297,8 +297,8 @@
     "\n",
     "model = Linear(1, 3, rngs=nnx.Rngs(0))\n",
     "\n",
-    "print(f\"{nnx.vars_as(model, mutable=False) = !s}\")\n",
-    "print(f\"{nnx.vars_as(model, mutable=True) = !s}\")"
+    "print(f\"{nnx.with_vars(model, mutable=False) = !s}\")\n",
+    "print(f\"{nnx.with_vars(model, mutable=True) = !s}\")"
    ]
   },
   {
@@ -317,7 +317,7 @@
    ],
    "source": [
     "v = nnx.Variable(jnp.array(0))\n",
-    "v_immut = nnx.vars_as(v, mutable=False)\n",
+    "v_immut = nnx.with_vars(v, mutable=False)\n",
     "assert not v_immut.mutable\n",
     "\n",
     "try:\n",
@@ -355,7 +355,7 @@
    ],
    "source": [
     "v = nnx.Variable(jnp.array(0))\n",
-    "v_ref = nnx.vars_as(v, ref=True)\n",
+    "v_ref = nnx.with_vars(v, ref=True)\n",
     "assert v_ref.ref\n",
     "print(v_ref)\n",
     "print(v_ref.get_raw_value())"
@@ -386,11 +386,11 @@
     }
    ],
    "source": [
-    "v_immut = nnx.vars_as(v_ref, mutable=False)\n",
+    "v_immut = nnx.with_vars(v_ref, mutable=False)\n",
     "assert not v_immut.ref\n",
     "print(\"immutable =\", v_immut)\n",
     "\n",
-    "v_ref = nnx.vars_as(v_immut, mutable=True)\n",
+    "v_ref = nnx.with_vars(v_immut, mutable=True)\n",
     "assert v_ref.ref\n",
     "print(\"mutable =\", v_ref)"
    ]
@@ -458,7 +458,7 @@
     "    model =  nnx.merge(graphdef, params, nondiff)\n",
     "    return ((model(x) - y) ** 2).mean()\n",
     "\n",
-    "  loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, mutable=False))  # immutable for jax.grad\n",
+    "  loss, grads = jax.value_and_grad(loss_fn)(nnx.with_vars(params, mutable=False))  # immutable for jax.grad\n",
     "  optimizer.update(model, grads)\n",
     "\n",
     "  return loss\n",
@@ -563,9 +563,9 @@
    "source": [
     "@jax.jit\n",
     "def create_model(rngs):\n",
-    "  return nnx.vars_as((Block(2, 64, 3, rngs=rngs)), hijax=False)\n",
+    "  return nnx.with_vars((Block(2, 64, 3, rngs=rngs)), hijax=False)\n",
     "\n",
-    "model = nnx.vars_as(create_model(nnx.Rngs(0)), hijax=True)\n",
+    "model = nnx.with_vars(create_model(nnx.Rngs(0)), hijax=True)\n",
     "\n",
     "print(\"model.linear =\", model.linear)"
    ]

docs_nnx/hijax/hijax.md

@@ -29,7 +29,7 @@ optimizer = nnx.Optimizer(model, optax.adamw(1e-2), wrt=nnx.Param)
 @jax.jit
 def train_step(x, y):
   loss_fn = lambda m: jnp.mean((m(x) - y) ** 2)
-  loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(model, mutable=False))  # tmp fix for jax.grad
+  loss, grads = jax.value_and_grad(loss_fn)(nnx.with_vars(model, mutable=False))  # tmp fix for jax.grad
   optimizer.update(model, grads)
   return loss
 
@@ -112,13 +112,13 @@ class Linear(nnx.Module):
 
 model = Linear(1, 3, rngs=nnx.Rngs(0))
 
-print(f"{nnx.vars_as(model, mutable=False) = !s}")
-print(f"{nnx.vars_as(model, mutable=True) = !s}")
+print(f"{nnx.with_vars(model, mutable=False) = !s}")
+print(f"{nnx.with_vars(model, mutable=True) = !s}")
 ```
 
 ```{code-cell} ipython3
 v = nnx.Variable(jnp.array(0))
-v_immut = nnx.vars_as(v, mutable=False)
+v_immut = nnx.with_vars(v, mutable=False)
 assert not v_immut.mutable
 
 try:
@@ -131,18 +131,18 @@ except Exception as e:
 
 ```{code-cell} ipython3
 v = nnx.Variable(jnp.array(0))
-v_ref = nnx.vars_as(v, ref=True)
+v_ref = nnx.with_vars(v, ref=True)
 assert v_ref.ref
 print(v_ref)
 print(v_ref.get_raw_value())
 ```
 
 ```{code-cell} ipython3
-v_immut = nnx.vars_as(v_ref, mutable=False)
+v_immut = nnx.with_vars(v_ref, mutable=False)
 assert not v_immut.ref
 print("immutable =", v_immut)
 
-v_ref = nnx.vars_as(v_immut, mutable=True)
+v_ref = nnx.with_vars(v_immut, mutable=True)
 assert v_ref.ref
 print("mutable =", v_ref)
 ```
@@ -176,7 +176,7 @@ def train_step(model, optimizer, x, y):
     model =  nnx.merge(graphdef, params, nondiff)
     return ((model(x) - y) ** 2).mean()
 
-  loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, mutable=False))  # immutable for jax.grad
+  loss, grads = jax.value_and_grad(loss_fn)(nnx.with_vars(params, mutable=False))  # immutable for jax.grad
   optimizer.update(model, grads)
 
   return loss
@@ -226,9 +226,9 @@ except Exception as e:
 ```{code-cell} ipython3
 @jax.jit
 def create_model(rngs):
-  return nnx.vars_as((Block(2, 64, 3, rngs=rngs)), hijax=False)
+  return nnx.with_vars((Block(2, 64, 3, rngs=rngs)), hijax=False)
 
-model = nnx.vars_as(create_model(nnx.Rngs(0)), hijax=True)
+model = nnx.with_vars(create_model(nnx.Rngs(0)), hijax=True)
 
 print("model.linear =", model.linear)
 ```

docs_nnx/mnist_tutorial.ipynb

@@ -303,7 +303,7 @@
     "\n",
     "## 6. Train and evaluate the model\n",
     "\n",
-    "Now, you can train the CNN model. Before the training loop, we use [`nnx.view`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation."
+    "Now, you can train the CNN model. Before the training loop, we use [`nnx.with_modules`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation."
    ]
   },
   {
@@ -335,8 +335,8 @@
     "}\n",
     "\n",
     "rngs = nnx.Rngs(0)\n",
-    "train_model = nnx.view(model, deterministic=False, use_running_average=False)\n",
-    "eval_model = nnx.view(model, deterministic=True, use_running_average=True)\n",
+    "train_model = nnx.with_modules(model, deterministic=False, use_running_average=False)\n",
+    "eval_model = nnx.with_modules(model, deterministic=True, use_running_average=True)\n",
     "\n",
     "for step, batch in enumerate(train_ds.as_numpy_iterator()):\n",
     "  # Run the optimization for one step and make a stateful update to the following:\n",
@@ -380,7 +380,7 @@
    "source": [
     "## 7. Perform inference on the test set\n",
     "\n",
-    "Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (an `nnx.view` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance."
+    "Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (using `nnx.with_modules` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance."
    ]
   },
   {

docs_nnx/mnist_tutorial.md

@@ -173,7 +173,7 @@ In the code above, the [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_ref
 
 ## 6. Train and evaluate the model
 
-Now, you can train the CNN model. Before the training loop, we use [`nnx.view`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation.
+Now, you can train the CNN model. Before the training loop, we use [`nnx.with_modules`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation.
 
 ```{code-cell} ipython3
 from IPython.display import clear_output
@@ -187,8 +187,8 @@ metrics_history = {
 }
 
 rngs = nnx.Rngs(0)
-train_model = nnx.view(model, deterministic=False, use_running_average=False)
-eval_model = nnx.view(model, deterministic=True, use_running_average=True)
+train_model = nnx.with_modules(model, deterministic=False, use_running_average=False)
+eval_model = nnx.with_modules(model, deterministic=True, use_running_average=True)
 
 for step, batch in enumerate(train_ds.as_numpy_iterator()):
   # Run the optimization for one step and make a stateful update to the following:
@@ -227,7 +227,7 @@ for step, batch in enumerate(train_ds.as_numpy_iterator()):
 
 ## 7. Perform inference on the test set
 
-Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (an `nnx.view` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance.
+Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (using `nnx.with_modules` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance.
 
 ```{code-cell} ipython3
 @nnx.jit

examples/nnx_toy_examples/hijax_basic.py

@@ -68,7 +68,7 @@ def loss_fn(params):
     model = nnx.merge(graphdef, params, nondiff)
     return jnp.mean((y - model(x)) ** 2)
 
-  grads = jax.grad(loss_fn)(nnx.vars_as(params, is_mutable=False))
+  grads = jax.grad(loss_fn)(nnx.with_vars(params, is_mutable=False))
   optimizer.update(model, grads)
 
 @jax.jit

examples/nnx_toy_examples/hijax_demo.py

@@ -238,7 +238,7 @@ def loss_fn(params):
 
   # For the time being we have to use 'immutable'
   # as 'jax.grad' doesn't support QDD types yet.
-  grads = jax.grad(loss_fn)(nnx.vars_as(params, is_mutable=False))
+  grads = jax.grad(loss_fn)(nnx.with_vars(params, is_mutable=False))
   # 'update' mutates the optimizer's state and the params in place
   # so we don't need to return anything 🚀
   optimizer.update(params, grads)

flax/nnx/__init__.py

@@ -50,7 +50,7 @@
 from .module import M as M
 from .module import Module as Module
 from .module import capture as capture
-from .module import view as view
+from .module import with_modules as with_modules
 from .module import view_info as view_info
 from .module import with_attributes as with_attributes
 from .module import iter_children as iter_children, iter_modules as iter_modules
@@ -75,8 +75,8 @@
 from .graphlib import MergeContext as MergeContext
 from .graphlib import merge_context as merge_context
 from .graphlib import variables as variables
-from .graphlib import vars_as as vars_as
-from .graphlib import pure as pure
+from .graphlib import with_vars as with_vars
+from .graphlib import as_pure as as_pure
 from .graphlib import cached_partial as cached_partial
 from .graphlib import flatten as flatten
 from .graphlib import unflatten as unflatten
@@ -152,7 +152,7 @@
 from .spmd import get_named_sharding as get_named_sharding
 from .spmd import with_partitioning as with_partitioning
 from .spmd import get_abstract_model as get_abstract_model
-from .spmd import abstract_with_sharding as abstract_with_sharding
+from .spmd import as_abstract as as_abstract
 from .statelib import FlatState as FlatState
 from .statelib import State as State
 from .statelib import to_flat_state as to_flat_state

flax/nnx/compat.py

@@ -39,7 +39,7 @@
 recursive_map = functools.partial(_graphlib.recursive_map, graph=True)
 
 # module
-view = functools.partial(_module.view, graph=True)
+view = functools.partial(_module.with_modules, graph=True)
 view_info = functools.partial(_module.view_info, graph=True)
 iter_modules = functools.partial(_module.iter_modules, graph=True)
 iter_children = functools.partial(_module.iter_children, graph=True)  # type: ignore[has-type]