Merge pull request #4766 from IvyZX:bdg

PiperOrigin-RevId: 771171776

Author: Flax Authors

flax/nnx/bridge/wrappers.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import dataclasses
+from functools import partial
 import typing as tp
 from typing import Any
 
@@ -133,6 +134,16 @@ def lazy_init(self, *args, **kwargs):
     """A shortcut of calling `nnx.bridge.lazy_init()` upon this module."""
     return lazy_init(self, *args, **kwargs)
 
+  def __getattr__(self, name: str):
+    if hasattr(super(), name):
+      return super().__getattribute__(name)
+    maybe_method = getattr(self.module.__class__, name, None)
+    if callable(maybe_method):
+      method = partial(self.__call__, method=maybe_method)
+      method.__self__ = self
+      return method
+    return super().__getattribute__(name)
+
   def __call__(
     self, *args: Any, rngs: tp.Optional[Rngs] = None,
     method: tp.Callable[..., Any] | str | None = None, **kwargs: Any

tests/nnx/bridge/wrappers_test.py

@@ -97,14 +97,18 @@ def dot(self, x):
         w = self.param('w', nn.initializers.lecun_normal(), (4, 3))
         return x @ w
 
+      def rngs(self):
+        raise ValueError('This should not be called because ToNNX has .rngs')
+
     x = jax.random.normal(jax.random.key(0), (2, 4))
     model = bridge.ToNNX(Foo(), rngs=nnx.Rngs(0))
-    bridge.lazy_init(model, x, method=model.module.dot)
-    y = model(x, method=model.module.dot)
+    bridge.lazy_init(model.dot, x)
+    y = model.dot(x)
     np.testing.assert_allclose(y, x @ nnx.state(model)['w'].value)
     # lazy_init only initialized param w inside dot(), so calling __call__ should fail
     with self.assertRaises(flax.errors.ScopeParamNotFoundError):
       y = model(x)
+    assert isinstance(model.rngs, nnx.Rngs)
 
   def test_linen_to_nnx_mutable(self):
     class Foo(nn.Module):