Merge pull request #2025 from jheek:fix-tranform-state-reuse

PiperOrigin-RevId: 439568379

Author: Flax Authors

flax/linen/transforms.py

@@ -290,6 +290,7 @@ def create_trans_fn(fn_name, fn_trafo_args):
     # we need to create a scope-function from our class for the given method
     @functools.wraps(fn)
     def wrapped_fn(self, *args, **kwargs):
+      state = self._state.export()
       # make a scope-function to transform
       def core_fn(scopes, *args, **kwargs):
         # make a clone of self using its arguments
@@ -301,7 +302,7 @@ def core_fn(scopes, *args, **kwargs):
         # we reference module_class, not self.__class__ to avoid infinite loop
         cloned = module_class(parent=None, **attrs)
         cloned, args, kwargs = set_module_scopes(cloned, args, kwargs, scopes)
-        object.__setattr__(cloned, '_state', self._state.export())  # pylint: disable=protected-access
+        object.__setattr__(cloned, '_state', state.export())  # pylint: disable=protected-access
         res = fn(cloned, *args, **kwargs)
         self._state.reimport(cloned._state)  # pylint: disable=protected-access
         _test_transformed_return_values(res, fn_name)
@@ -343,12 +344,13 @@ def decorator_lift_transform(transform, class_fn, *trafo_args,
   prewrapped_fns = [wrap_method_once(class_fn) for class_fn in class_fns]
   @functools.wraps(prewrapped_fns[0])
   def wrapped_fn(self, *args, **kwargs):
+    state = self._state.export()
     # make a scope-function to transform
     def core_fn(prewrapped_fn, class_fn, scopes, *args, **kwargs):
       if not multi_scope:
         scopes = [scopes]
       cloned, args, kwargs = set_module_scopes(self, args, kwargs, scopes)
-      object.__setattr__(cloned, '_state', self._state.export())  # pylint: disable=protected-access
+      object.__setattr__(cloned, '_state', state.export())  # pylint: disable=protected-access
       res = prewrapped_fn(cloned, *args, **kwargs)
       self._state.reimport(cloned._state)  # pylint: disable=protected-access
       _test_transformed_return_values(res, getattr(class_fn, '__name__', None))

tests/linen/linen_transforms_test.py

@@ -1434,6 +1434,24 @@ def __call__(self, x):
     with self.assertRaises(errors.TransformTargetError):
       Foo().init(random.PRNGKey(0), jnp.zeros((2, 3)))
 
+  def test_scan_compact_count(self):
+    class Foo(nn.Module):
+      num_layers: int = 5
+
+      @nn.compact
+      def __call__(self, x):
+        def body_fn(mdl, x):
+          return nn.Dense(features=x.shape[-1])(x), ()
+        x, _ = nn.scan(body_fn, length=self.num_layers, variable_axes={"params": 0}, split_rngs={"params": True})(self, x)
+        return x
+
+    m = Foo()
+    x = jnp.ones((3,))
+    v = m.init(jax.random.PRNGKey(0), x)
+    self.assertEqual(v['params']['Dense_0']['kernel'].shape, (5, 3, 3))
+    m.apply(v, x)
+
+
 
 if __name__ == '__main__':
   absltest.main()