deprecate Variable.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
+   @nnx.jit  # automatic state propagation
    def train_step(model, optimizer, x, y):
-     def loss_fn(model):
-       y_pred = model(x)  # call methods directly
-       return ((y_pred - y) ** 2).mean()
-
+     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
 
 

examples/gemma/helpers.py

@@ -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_test.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/layers.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/modules.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/sampler_test.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/sow_lib.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/transformer.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
 
 

examples/gemma/transformer_test.py

@@ -461,7 +461,7 @@ def test_sow_intermediates(self, sow_config):
 
     if sow_config.embeddings:
       self.assertTrue(hasattr(transformer, 'embeddings'))
-      embeddings = transformer.embeddings.value[0]
+      embeddings = transformer.embeddings[0]
       self.assertEqual(
           embeddings.shape,
           (batch_size, sequence_length, config.embed_dim),
@@ -472,7 +472,7 @@ def test_sow_intermediates(self, sow_config):
     for layer in transformer.layers:
       if sow_config.rs_after_attention:
         self.assertTrue(hasattr(layer, 'rs_after_attention'))
-        rs_after_attention = layer.rs_after_attention.value[0]
+        rs_after_attention = layer.rs_after_attention[0]
         self.assertIsNotNone(rs_after_attention)
         self.assertEqual(
             rs_after_attention.shape,
@@ -482,7 +482,7 @@ def test_sow_intermediates(self, sow_config):
         self.assertFalse(hasattr(layer, 'rs_after_attention'))
       if sow_config.rs_after_ffw:
         self.assertTrue(hasattr(layer, 'rs_after_ffw'))
-        rs_after_ffw = layer.rs_after_ffw.value[0]
+        rs_after_ffw = layer.rs_after_ffw[0]
         self.assertIsNotNone(rs_after_ffw)
         self.assertEqual(
             rs_after_ffw.shape,
@@ -492,7 +492,7 @@ def test_sow_intermediates(self, sow_config):
         self.assertFalse(hasattr(layer, 'rs_after_ffw'))
       if sow_config.attn_logits_topk:
         self.assertTrue(hasattr(layer.attn, 'logits_topk_values'))
-        attn_logits_topk_values = layer.attn.logits_topk_values.value[0]
+        attn_logits_topk_values = layer.attn.logits_topk_values[0]
         self.assertIsNotNone(attn_logits_topk_values)
         self.assertEqual(
             attn_logits_topk_values.shape,
@@ -504,7 +504,7 @@ def test_sow_intermediates(self, sow_config):
             ),
         )
         self.assertTrue(hasattr(layer.attn, 'logits_topk_indices'))
-        attn_logits_topk_indices = layer.attn.logits_topk_indices.value[0]
+        attn_logits_topk_indices = layer.attn.logits_topk_indices[0]
         self.assertIsNotNone(attn_logits_topk_indices)
         self.assertEqual(
             attn_logits_topk_indices.shape,
@@ -520,7 +520,7 @@ def test_sow_intermediates(self, sow_config):
         self.assertFalse(hasattr(layer.attn, 'logits_topk_indices'))
       if sow_config.mlp_hidden_topk:
         self.assertTrue(hasattr(layer.mlp, 'hidden_topk_values'))
-        ffw_hidden_topk_values = layer.mlp.hidden_topk_values.value[0]
+        ffw_hidden_topk_values = layer.mlp.hidden_topk_values[0]
         self.assertIsNotNone(ffw_hidden_topk_values)
         self.assertEqual(
             ffw_hidden_topk_values.shape,
@@ -531,7 +531,7 @@ def test_sow_intermediates(self, sow_config):
             ),
         )
         self.assertTrue(hasattr(layer.mlp, 'hidden_topk_indices'))
-        ffw_hidden_topk_indices = layer.mlp.hidden_topk_indices.value[0]
+        ffw_hidden_topk_indices = layer.mlp.hidden_topk_indices[0]
         self.assertIsNotNone(ffw_hidden_topk_indices)
         self.assertEqual(
             ffw_hidden_topk_indices.shape,

flax/configurations.py

@@ -201,6 +201,38 @@ def static_bool_env(varname: str, default: bool) -> bool:
     )
 
 
+def str_flag(name: str, *, default: str, help: str) -> FlagHolder[str]:
+  """Set up a string flag.
+
+  Example::
+
+    some_string = str_flag(
+        name='flax_some_string',
+        default='default_value',
+        help='Some string configuration.',
+    )
+
+  Now the ``FLAX_SOME_STRING`` shell environment variable can be used to
+  control the process-level value of the flag, in addition to using e.g.
+  ``config.update("flax_some_string", "new_value")`` directly.
+
+  Args:
+    name: converted to lowercase to define the name of the flag. It is
+      converted to uppercase to define the corresponding shell environment
+      variable.
+    default: a default value for the flag.
+    help: used to populate the docstring of the returned flag holder object.
+
+  Returns:
+    A flag holder object for accessing the value of the flag.
+  """
+  name = name.lower()
+  config._add_option(name, static_str_env(name.upper(), default))
+  fh = FlagHolder[str](name, help)
+  setattr(Config, name, property(lambda _: fh.value, doc=help))
+  return fh
+
+
 def static_int_env(varname: str, default: int | None) -> int | None:
   """Read an environment variable and interpret it as an integer.
 
@@ -222,6 +254,18 @@ def static_int_env(varname: str, default: int | None) -> int | None:
     ) from None
 
 
+def static_str_env(varname: str, default: str) -> str:
+  """Read an environment variable and interpret it as a string.
+
+  Args:
+    varname: the name of the variable
+    default: the default string value
+  Returns:
+    string return value derived from defaults and environment.
+  """
+  return os.getenv(varname, default)
+
+
 # Flax Global Configuration Variables:
 
 flax_filter_frames = bool_flag(
@@ -294,5 +338,5 @@ def static_int_env(varname: str, default: int | None) -> int | None:
 flax_hijax_variable = bool_flag(
   name='flax_hijax_variable',
   default=False,
-  help='Whether to enable HiJAX support for `nnx.Variable`.',
+  help='Whether to use hijax for `nnx.Variable`. Options are "pytree", "hijax", and "ref".',
 )