deprecate .value

Author: cgarciae

docs_nnx/index.rst

@@ -107,15 +107,11 @@ Basic usage
    model = Model(2, 64, 3, rngs=nnx.Rngs(0))  # eager initialization
    optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
 
-   @nnx.jit  # automatic state management for JAX transforms
-   def train_step(model, optimizer, x, y, rngs):
-     def loss_fn(model):
-       y_pred = model(x, rngs)  # call methods directly
-       return ((y_pred - y) ** 2).mean()
-
+   @nnx.jit  # automatic state propagation
+   def train_step(model, optimizer, x, y):
+     loss_fn = lambda model: ((model(x) - y) ** 2).mean()
      loss, grads = nnx.value_and_grad(loss_fn)(model)
      optimizer.update(model, grads)  # in-place updates
-
      return loss
 
 

docs_nnx/nnx_basics.ipynb

@@ -40,7 +40,7 @@ class Linear(nnx.Module):
     self.din, self.dout = din, dout
 
   def __call__(self, x: jax.Array):
-    return x @ self.w + self.b
+    return x @ self.w + self.b[None]
 ```
 
 Also note that the inner values of `Variable`s can be accessed using the `value` property, but for convenience they implement all numeric operators and can be used directly in arithmetic expressions (as shown in the code above).
@@ -73,12 +73,12 @@ class Counter(nnx.Module):
     self.count = Count(jnp.array(0))
 
   def __call__(self):
-    self.count.value += 1
+    self.count[...] += 1
 
 counter = Counter()
-print(f'{counter.count.value = }')
+print(f'{counter.count[...] = }')
 counter()
-print(f'{counter.count.value = }')
+print(f'{counter.count[...] = }')
 ```
 
 Mutable references are usually avoided in JAX. But Flax NNX provides sound mechanisms

docs_nnx/nnx_basics.md

@@ -74,7 +74,7 @@ def assign_val_fn(
         mapped_path: tuple[str | int, ...],
         val: Any,
     ) -> dict[tuple[str, ...], Any]:
-      state[mapped_path].value = val
+      state[mapped_path].set_value(val)
       return state
 
   mdl: M = nnx.eval_shape(module_factory)

examples/gemma/helpers.py

@@ -137,11 +137,11 @@ def _map_key_fn(key: tuple[str, ...]) -> tuple[str | int, ...]:
     np.testing.assert_array_equal(output, linen_output)
     for i in range(len(num_features)):
       np.testing.assert_array_equal(
-          mdl.layers[i].layers[0].mean.value,
+          mdl.layers[i].layers[0].mean[...],
           linen_vars['batch_stats'][f'layers_{i}']['layers_0']['mean'],
       )
       np.testing.assert_array_equal(
-          mdl.layers[i].layers[0].var.value,
+          mdl.layers[i].layers[0].var[...],
           linen_vars['batch_stats'][f'layers_{i}']['layers_0']['var'],
       )
 

examples/gemma/helpers_test.py

@@ -44,11 +44,11 @@ def __init__(
     self.w = nnx.Param(kernel_init(rngs.params(), shape, dtype))
 
   def __call__(self, x: ArrayLike) -> Array:
-    return jnp.einsum(self.einsum_str, x, self.w.value)
+    return jnp.einsum(self.einsum_str, x, self.w[...])
 
   @property
   def shape(self) -> Shape:
-    return self.w.value.shape
+    return self.w.shape
 
 
 class RMSNorm(nnx.Module):
@@ -65,12 +65,12 @@ def __init__(
     self.scale = nnx.Param(scale_init(rngs.params(), dim, dtype))
 
   def __call__(self, x: Array) -> Array:
-    dtype = self.scale.value.dtype
+    dtype = self.scale.dtype
     var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
     normed_inputs = jnp.asarray(x * jax.lax.rsqrt(var + 1e-06), dtype=dtype)
     # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is
     # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to
     # a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs.
-    scale = jnp.expand_dims(self.scale.value, axis=range(len(x.shape) - 1))
+    scale = jnp.expand_dims(self.scale, axis=range(len(x.shape) - 1))
     normed_inputs = normed_inputs * (1 + scale)
     return normed_inputs

examples/gemma/layers.py

@@ -63,15 +63,15 @@ def encode(self, x: ArrayLike) -> Array:
     return x
 
   def decode(self, x: ArrayLike) -> Array:
-    return jnp.dot(x, self.input_embedding.value.T)
+    return jnp.dot(x, self.input_embedding.T)
 
   @property
   def embed_dim(self):
-    return self.input_embedding.value.shape[1]
+    return self.input_embedding.shape[1]
 
   @property
   def num_embed(self):
-    return self.input_embedding.value.shape[0]
+    return self.input_embedding.shape[0]
 
 
 class Attention(nnx.Module):

examples/gemma/modules.py

@@ -232,9 +232,9 @@ def test_forbidden_tokens(self):
         transformer_config, rngs=nnx.Rngs(params=0)
     )
     # Pre-cook the embedding matrix so that the output is deterministic.
-    transformer.embedder.input_embedding.value = jnp.eye(
+    transformer.embedder.input_embedding.set_value(jnp.eye(
         vocab.GetPieceSize(), 32
-    )
+    ))
     sampler = sampler_lib.Sampler(
         transformer=transformer,
         vocab=vocab,

examples/gemma/sampler_test.py

@@ -49,13 +49,11 @@ def merge(self, decoding_step, layer: nnx.Module):
         if field.name.startswith('attn_'):
           step_value = getattr(
               layer.attn, field.name.replace('attn_', '')
-          ).value[0]
+          )[0]
         elif field.name.startswith('mlp_'):
-          step_value = getattr(layer.mlp, field.name.replace('mlp_', '')).value[
-              0
-          ]
+          step_value = getattr(layer.mlp, field.name.replace('mlp_', ''))[0]
         else:
-          step_value = getattr(layer, field.name).value[0]
+          step_value = getattr(layer, field.name)[0]
       except AttributeError as exc:
         raise ValueError(
             f'Intermediate {field.name} is not in the step intermediates.'
@@ -93,7 +91,7 @@ def merge(self, decoding_step, transformer: nnx.Module):
     if self.embeddings is not None:
       try:
         self.embeddings = self.embeddings.at[:, decoding_step + 1, ...].set(
-            transformer.embeddings.value[0][:, 0, ...]
+            transformer.embeddings[0][:, 0, ...]
         )
       except AttributeError as exc:
         raise ValueError(

examples/gemma/sow_lib.py

@@ -487,10 +487,10 @@ def _assign_linen_params_to_nnx_state(
   if 'gate_proj' in mapped_path:
     if transpose_gating_einsum:
       val = jnp.swapaxes(val, 1, 2)
-    state[mapped_path].value = val[0]
-    state[mapped_path[:-2] + ('up_proj', 'kernel')].value = val[1]
+    state[mapped_path].set_value(val[0])
+    state[mapped_path[:-2] + ('up_proj', 'kernel')].set_value(val[1])
   else:
-    state[mapped_path].value = val
+    state[mapped_path].set_value(val)
   return state