Deprecate Variable hooks in favor of custom Variable types.

flax/nnx/variablelib.py

@@ -1354,6 +1354,13 @@ def __init__(
 
     _remap_sharding_metadata(metadata)
 
+    hook_keys = ['on_get_value', 'on_set_value', 'on_create_value', 'on_add_axis', 'on_remove_axis']
+    if any(k in metadata and metadata[k] for k in hook_keys):
+      warnings.warn(
+        'Variable hooks are deprecated. Create your own Variable type and '
+        'overload `get_value` and `set_value` instead.'
+      )
+
     # run create_value hooks
     if 'on_create_value' in metadata:
       hooks = metadata['on_create_value']
@@ -2259,6 +2266,12 @@ def with_metadata(
   ] = (),
   **metadata: tp.Any,
 ) -> F:
+  if on_set_value or on_get_value or on_create_value or on_add_axis or on_remove_axis:
+    warnings.warn(
+      'Variable hooks are deprecated in favor of users creating their own '
+      'Variable types and overloading `get_value` and `set_value` instead.'
+    )
+
   if on_set_value:
     if callable(on_set_value):
       on_set_value = (on_set_value,)

tests/nnx/containers_test.py

@@ -20,39 +20,46 @@
 
 class TestContainers(absltest.TestCase):
   def test_unbox(self):
-    x = nnx.Param(
-      jnp.array(1),
-      on_get_value=lambda c, x: x + 3,  # type: ignore
-    )
+    class CustomParam(nnx.Param):
+      def get_value(self, **kwargs):
+        return super().get_value(**kwargs) + 3
+
+    x = CustomParam(jnp.array(1))
 
     assert x[...] == 4
 
   def test_on_set_value(self):
-    x = nnx.Param(
-      jnp.array(1),  # type: ignore
-      on_set_value=lambda c, x: x + 7,  # type: ignore
-    )
+    class CustomParam(nnx.Param):
+      def set_value(self, value, **kwargs):
+        super().set_value(value + 7, **kwargs)
+
+    x = CustomParam(jnp.array(1))
     x[...] = 5
 
     assert x.get_raw_value() == 12
 
   def test_module_unbox(self):
+    class CustomParam(nnx.Param):
+      def get_value(self, **kwargs):
+        return super().get_value(**kwargs) + 3
+
     class Foo(nnx.Module):
       def __init__(self) -> None:
-        self.x = nnx.Param(1, on_get_value=lambda c, x: x + 3)
+        self.x = CustomParam(1)
 
     module = Foo()
 
     assert module.x.get_value() == 4
     assert vars(module)['x'].get_raw_value() == 1
 
   def test_module_box(self):
+    class CustomParam(nnx.Param):
+      def set_value(self, value, **kwargs):
+        super().set_value(value + 7, **kwargs)
+
     class Foo(nnx.Module):
       def __init__(self) -> None:
-        self.x = nnx.Param(
-          jnp.array(1),
-          on_set_value=lambda c, x: x + 7,  # type: ignore
-        )
+        self.x = CustomParam(jnp.array(1))
 
     module = Foo()
     module.x[...] = 5

tests/nnx/spmd_test.py

@@ -150,6 +150,24 @@ def __call__(self, x):
   def test_add_remove_axis_in_transform(self):
     test = self
     kadds, kremoves, badds, bremoves = [], [], [], []
+    class CustomVar(nnx.Param):
+      def add_axis(self, axis_index, axis_name):
+        kadds.append((axis_index, axis_name))
+        super().add_axis(axis_index, axis_name)
+
+      def remove_axis(self, axis_index, axis_name):
+        kremoves.append((axis_index, axis_name))
+        super().remove_axis(axis_index, axis_name)
+
+    class CustomBiasVar(nnx.Param):
+      def add_axis(self, axis_index, axis_name):
+        badds.append((axis_index, axis_name))
+        super().add_axis(axis_index, axis_name)
+
+      def remove_axis(self, axis_index, axis_name):
+        bremoves.append((axis_index, axis_name))
+        super().remove_axis(axis_index, axis_name)
+
     class MLP(nnx.Module):
 
       @nnx.split_rngs(splits=5, graph=True, graph_updates=True)
@@ -159,22 +177,13 @@ class MLP(nnx.Module):
           graph=True, graph_updates=True,
       )
       def __init__(self, rngs: nnx.Rngs):
-        self.linear = nnx.Linear(
-          4,
-          4,
-          kernel_init=nnx.with_metadata(
-            nnx.initializers.lecun_normal(),
+        self.kernel = CustomVar(
+            nnx.initializers.lecun_normal()(rngs.params(), (4, 4)),
             out_sharding=('din', 'dout'),
             nickname=('in', 'out'),
-            on_add_axis=lambda _, idx, name: kadds.append((idx, name)),
-            on_remove_axis=lambda _, idx, name: kremoves.append((idx, name)),
-          ),
-          bias_init=nnx.with_metadata(
-            nnx.initializers.zeros_init(),  # no sharding annotation here!
-            on_add_axis=lambda _, idx, name: badds.append((idx, name)),
-            on_remove_axis=lambda _, idx, name: bremoves.append((idx, name)),
-          ),
-          rngs=rngs,
+        )
+        self.bias = CustomBiasVar(
+            nnx.initializers.zeros_init()(rngs.params(), (4,))
         )
 
       @nnx.scan(
@@ -183,10 +192,10 @@ def __init__(self, rngs: nnx.Rngs):
           graph=True, graph_updates=True,
       )
       def __call__(self, x: jax.Array):
-        x = self.linear(x)
+        x = x @ self.kernel + self.bias[None, :]
         # test sharding layer axes is not present inside scan
-        test.assertEqual(self.linear.kernel.shape, (4, 4))
-        test.assertEqual(self.linear.kernel.out_sharding, ('din', 'dout'))
+        test.assertEqual(self.kernel.shape, (4, 4))
+        test.assertEqual(self.kernel.out_sharding, ('din', 'dout'))
         # at least a remove_axis was already called to remove the layer axis
         test.assertEqual(kremoves[-1], (0, 'layers'))
         test.assertEqual(bremoves[-1], (0, 'layers'))
@@ -200,10 +209,10 @@ def __call__(self, x: jax.Array):
     )
     with jax.set_mesh(mesh):
       m = MLP(rngs=nnx.Rngs(0))
-    self.assertEqual(m.linear.kernel.shape, (5, 4, 4))
-    self.assertEqual(m.linear.kernel.out_sharding, ('layers', 'din', 'dout'))
-    self.assertEqual(m.linear.kernel.nickname, ('nick', 'in', 'out'))
-    self.assertEqual(m.linear.bias.shape, (5, 4))
+    self.assertEqual(m.kernel.shape, (5, 4, 4))
+    self.assertEqual(m.kernel.out_sharding, ('layers', 'din', 'dout'))
+    self.assertEqual(m.kernel.nickname, ('nick', 'in', 'out'))
+    self.assertEqual(m.bias.shape, (5, 4))
     # One add_axis called to add the `nnx.vmap` dimension
     self.assertEqual(kadds, [(0, 'layers')])
     self.assertEqual(kremoves, [])
@@ -242,7 +251,6 @@ def f(v):
     self.assertEqual(v2.out_sharding, ('dmid', 'din', 'dout'))
     self.assertEqual(v2[...], 10)
 
-
   def test_transform_metadata_decorator_none_partition(self):
     v = nnx.Param(
       jnp.array(1),
@@ -437,8 +445,8 @@ def test_eval_shape_with_explicit_sharding(self):
     axis_types = (jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit)
     mesh1 = jax.make_mesh((2, 2), ("a", "b"), axis_types)
     class Model(nnx.Module):
-        def __init__(self):
-          self.p1 = nnx.Param(
+      def __init__(self):
+        self.p1 = nnx.Param(
             reshard(jnp.ones((4,4)), NamedSharding(mesh1, P('a', 'b'))),
             mesh=mesh1)
 
@@ -451,9 +459,9 @@ def test_eval_shape_with_sharding0(self):
     mesh2 = jax.make_mesh((1, 4), ("c", "d"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))
 
     class Model(nnx.Module):
-        def __init__(self):
-            self.p1 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b"), "mesh": mesh1})
-            self.p2 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("c", "d"), "mesh": mesh2})
+      def __init__(self):
+        self.p1 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b"), "mesh": mesh1})
+        self.p2 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("c", "d"), "mesh": mesh2})
 
     abs_model = nnx.as_abstract(nnx.eval_shape(lambda: Model()))
     assert isinstance(abs_model.p1.kernel.sharding, jax.sharding.NamedSharding)
@@ -465,12 +473,12 @@ def __init__(self):
 
   def test_eval_shape_with_sharding1(self):
     class Model(nnx.Module):
-        def __init__(self):
-            self.linear = nnx.Linear(10, 10, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b")})
+      def __init__(self):
+        self.linear = nnx.Linear(10, 10, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b")})
 
     mesh = jax.make_mesh((2, 2), ("a", "b"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))
     with jax.set_mesh(mesh):
-        abs_model = nnx.as_abstract(nnx.eval_shape(lambda: Model()))
+      abs_model = nnx.as_abstract(nnx.eval_shape(lambda: Model()))
     assert isinstance(abs_model.linear.kernel.sharding, jax.sharding.NamedSharding)
     assert abs_model.linear.kernel.sharding.mesh.axis_names == mesh.axis_names
     assert abs_model.linear.kernel.sharding.spec == jax.P("a", "b")

tests/nnx/summary_test.py

@@ -279,24 +279,21 @@ def __call__(self, x):
     # We should see 3 calls per block, plus one overall call, minus the shared call
     self.assertEqual(sum([s.startswith("├─") for s in table.splitlines()]), 6)
 
-  def test_tabulate_with_variable_hooks(self):
-    """Test that tabulate works with Variables implementing hooks and custom metadata."""
+  def test_tabulate_with_custom_variable(self):
+    """Test that tabulate works with custom Variables and custom metadata."""
 
     class Custom:
       def __repr__(self):
         return "<CustomMetadata>"
 
-    class VarWithHooks(nnx.Variable):
-        def on_get_value(self, value):
-            return value
-
-        def on_set_value(self, value):
-            return value + 1.0
+    class CustomVar(nnx.Variable):
+        def set_value(self, value, **kwargs):
+            return super().set_value(value + 1.0, **kwargs)
 
     class Model(nnx.Module):
         def __init__(self):
-            # Variable with hooks
-            self.hooked_param = VarWithHooks(value=jnp.ones((2, 3)))
+            # Variable with custom behavior
+            self.hooked_param = CustomVar(value=jnp.ones((2, 3)))
             self.hooked_param.set_metadata('description', 'Custom parameter')
             self.hooked_param.set_metadata('trainable', True)
 
@@ -315,7 +312,6 @@ def __call__(self, x):
     # Verify table contains expected content
     self.assertIn('Model Summary', table_repr)
     self.assertIn('hooked_param', table_repr)
-    self.assertIn('on_set_value', table_repr)
     self.assertIn('<CustomMetadata>', table_repr)
 
     # Verify metadata is preserved in the module