Merge pull request #2028 from levskaya:getattribute

PiperOrigin-RevId: 439935122

Author: Flax Authors

flax/linen/module.py

@@ -714,6 +714,17 @@ def __getattr__(self, name: str) -> Any:
       raise AttributeError(
           f'"{self.__class__.__name__}" object has no attribute "{name}"')
 
+  def __getattribute__(self, name):
+    """Call setup() before accessing any submodule attributes."""
+    # NB: all code here is very "hot" and will be run very frequently.
+    if name in object.__getattribute__(self, '__dataclass_fields__'):
+      if (name != 'parent' and
+          object.__getattribute__(self, '__dataclass_fields__')[name].init and
+          isinstance(object.__getattribute__(self, name), Module)):
+        object.__getattribute__(self, '_try_setup')()
+    # always run original python __getattribute__
+    return object.__getattribute__(self, name)
+
   def __dir__(self) -> Iterable[str]:
     """Call setup() before listing attributes."""
     self._try_setup()

tests/linen/linen_module_test.py

@@ -1517,5 +1517,23 @@ def __call__(self):
     self.assertTrue(foo.apply({}, rngs={'bar': k}))
     self.assertFalse(foo.apply({}, rngs={'baz': k}))
 
+  def test_getattribute_triggers_setup(self):
+    class B(nn.Module):
+      def setup(self):
+        self.p1 = self.param('p1', lambda k: jnp.ones((2,)))
+      def fn1(self, x):
+        return self.p1 + x
+    class A(nn.Module):
+      b: nn.Module
+      def __call__(self, x):
+        return self.b.fn1(x)
+    a = A(b=B())
+    k = random.PRNGKey(0)
+    x = jnp.zeros((2,))
+    vs = nn.init(lambda a,x: a(x), a)(k, x)
+    y = nn.apply(lambda a,x: a.b.fn1(x), a)(vs, x)
+    np.testing.assert_array_equal(y, jnp.ones((2,)))
+
+
 if __name__ == '__main__':
   absltest.main()