diff --git a/flax/nnx/variablelib.py b/flax/nnx/variablelib.py index 3de713404..80021d5c4 100644 --- a/flax/nnx/variablelib.py +++ b/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,) diff --git a/tests/nnx/containers_test.py b/tests/nnx/containers_test.py index 7aa80cb73..4bd136db4 100644 --- a/tests/nnx/containers_test.py +++ b/tests/nnx/containers_test.py @@ -20,26 +20,32 @@ 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() @@ -47,12 +53,13 @@ def __init__(self) -> None: 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 diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py index e803235b5..e1dcfc0d2 100644 --- a/tests/nnx/spmd_test.py +++ b/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") diff --git a/tests/nnx/summary_test.py b/tests/nnx/summary_test.py index cb6d918a6..e5003f0f4 100644 --- a/tests/nnx/summary_test.py +++ 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 "" - 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('', table_repr) # Verify metadata is preserved in the module