Merge pull request #2499 from levskaya:lifetimefix

PiperOrigin-RevId: 479153741

Author: Flax Authors

flax/linen/module.py

@@ -523,6 +523,27 @@ def reimport(self, other: '_ModuleInternalState') -> None:
 _ParentType = Union[Type['Module'], Type[Scope], Type[_Sentinel], None]
 
 
+class ParentDescriptor:
+  """Wraps parent module references in weak refs.
+
+  This prevents reference cycles from forming via parent links which can lead
+  to accidental OOMs in eager mode due to slow garbage collection as well as
+  spurious tracer leaks during jit compilation.
+
+  Note: "descriptors" are the underlying python mechanism for implementing
+  dynamic @property decorators.  We need to use a raw descriptor instead of the
+  more common decorator in order to force that the appropriate getter/setter
+  logic applies in subclasses even after various dataclass transforms.
+  """
+  def __get__(self, obj, objtype=None):
+    parent = object.__getattribute__(obj, "_parent_ref")
+    return parent() if isinstance(parent, weakref.ReferenceType) else parent
+
+  def __set__(self, obj, value):
+    maybe_weak = weakref.ref(value) if isinstance(value, Module) else value
+    object.__setattr__(obj, "_parent_ref", maybe_weak)
+
+
 # Base Module definition.
 # -----------------------------------------------------------------------------
 
@@ -588,6 +609,8 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
     # Set empty class defaults.
     cls._state = _uninitialized_module_internal_state
     cls.scope: Optional[Scope] = None
+    # Handles weak referencing of parent Modules to prevent reference cycles.
+    cls.parent = ParentDescriptor()
 
   @classmethod
   def _customized_dataclass_transform(cls):

tests/linen/linen_module_test.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 import functools
+import gc
 import operator
 from typing import (Any, Callable, Generic, Mapping, NamedTuple, Sequence,
                     Tuple, TypeVar)
@@ -31,6 +32,7 @@
 from jax.nn import initializers
 import jax.numpy as jnp
 import numpy as np
+from unittest.mock import patch
 
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
@@ -1755,5 +1757,31 @@ def __call__(self):
     self.assertFalse(foo.apply({}))
 
 
+class LeakTests(absltest.TestCase):
+
+  def test_tracer_leaks(self):
+    model = nn.Sequential([nn.Dense(50)])
+
+    @jax.jit
+    @functools.partial(jax.vmap, in_axes=(0, None))
+    def sample_from_prior(rng, inp):
+      params = model.init(rng, np.zeros((10, 50)))
+      out = model.apply(params, inp)
+      del params
+      return out
+
+    # disable manual gc.collect call in jax leak checker
+    # so that we can test tracer leaks in ref-cycles.  This is a
+    # reasonable proxy for transiently leaked memory during
+    # eager execution.
+    with patch.object(gc, 'collect', return_value=0):
+      with jax.checking_leaks():
+        for i in range(5):
+          rngs = jax.random.split(jax.random.PRNGKey(23), 100)
+          out = sample_from_prior(rngs, np.ones((4, 50)))
+          out.block_until_ready()
+          del out, rngs
+
+
 if __name__ == '__main__':
   absltest.main()