Merge pull request #2313 from levskaya:idfix

PiperOrigin-RevId: 463175098

flax/core/scope.py

@@ -22,14 +22,15 @@
 from typing import (Any, Callable, Dict, Generic, Iterable, Mapping, Optional,
                     Sequence, Set, Tuple, TypeVar, Union)
 
-from . import tracers
+from flax.ids import uuid
 from flax import config
 from flax import errors
 from flax import struct
 from flax import traceback_util
 from .frozen_dict import freeze
 from .frozen_dict import FrozenDict
 from .frozen_dict import unfreeze
+from . import tracers
 import jax
 from jax import config as jax_config
 from jax import numpy as jnp
@@ -51,7 +52,6 @@
 # When conditioning on filters we require explicit boolean comparisons.
 # pylint: disable=g-bool-id-comparison
 
-
 @dataclasses.dataclass(frozen=True, eq=True)
 class DenyList:
   """DenyList represents an opt-out based mutability filter.
@@ -343,6 +343,7 @@ def __init__(self, scope: 'Scope', collection: str, name: str):
       collection: The collection of the variable (e.g., "params").
       name: The name of the variable (e.g., "dense").
     """
+    self._id = uuid()
     self.scope = scope
     self.collection = collection
     self.name = name

flax/ids.py

@@ -0,0 +1,57 @@
+# Copyright 2022 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""UUIDs for Flax internals."""
+
+import threading
+
+
+class UUIDManager:
+  """Globally unique counter-based id manager.
+
+  We need globally unique key ids for Module and Variable object instances
+  to preserve and recreate sharing-by-reference relationship when lifting
+  transforms and adopting outside Modules.
+  - Use of id() is unacceptable because these identifiers are literally
+    pointers which can be recycled, so we rely on a globally unique counter id
+    instead.
+  - We need to handle copy/deepcopy uniqueness via a wrapped type.
+  """
+  def __init__(self):
+    self._lock = threading.Lock()
+    self._id = 0
+
+  def __call__(self):
+    with self._lock:
+      self._id += 1
+      return FlaxId(self._id)
+
+uuid = UUIDManager()
+
+
+class FlaxId:
+  """Hashable wrapper for ids that handles uniqueness of copies."""
+  def __init__(self, rawid):
+    self.id = rawid
+  def __eq__(self, other):
+    return isinstance(other, FlaxId) and other.id == self.id
+  def __hash__(self):
+    return hash(self.id)
+  def __repr__(self):
+    return f"FlaxId({self.id})"
+  def __deepcopy__(self, memo):
+    del memo
+    return uuid()
+  def __copy__(self):
+    return uuid()

flax/linen/module.py

@@ -36,6 +36,7 @@
 from flax.core.scope import (  # pylint: disable=g-multiple-import
     CollectionFilter, DenyList, FrozenVariableDict, Variable, VariableDict,
     union_filters)
+from flax.ids import uuid
 from flax.linen import summary
 
 
@@ -730,6 +731,7 @@ def __post_init__(self) -> None:
     # initialization, attach this Module as a submodule of a parent, or bind
     # this Module at the top-level to variables and rngs.
 
+    object.__setattr__(self, '_id', uuid())
     object.__setattr__(self, '_state', _ModuleInternalState())
 
     # Typically we set the parent based on the dynamic module context.
@@ -827,14 +829,12 @@ def adopt_attr_modules(cache, queue, suffix, subvalue):
           # Module was passed from outside. It needs to be cloned.
           # Outside modules are named by attachment, not an outer name.
           object.__setattr__(subvalue, 'name', None)
-          key = id(subvalue)
+          # Preserve sharing-by-reference relationships during adoption
+          # via cache keyed on unique instance ids.
+          key = subvalue._id
           if key not in cache:
-            # since we use id() as key, we need to keep a reference to original
-            # subvalue to ensure it's lifetime is long enough for the entire
-            # model setup and the id() is not recycled.
-            # TODO(levskaya): consider switching to per-module UUIDs
-            cache[key] = (subvalue.clone(), subvalue)
-          subvalue = cache[key][0]
+            cache[key] = subvalue.clone()
+          subvalue = cache[key]
         if subvalue.name is None:
           object.__setattr__(subvalue, 'parent', self)
           object.__setattr__(subvalue, 'name', f'{name}{suffix}')

flax/linen/transforms.py

@@ -43,6 +43,7 @@
 
 traceback_util.register_exclusion(__file__)
 
+# pylint: disable=protected-access
 
 # Utils
 # -----------------------------------------------------------------------------
@@ -79,12 +80,12 @@ def wrapped_fn(x):
     nonlocal refs
     if isinstance(x, (VariablePlaceholder, InstancePlaceholder)):
       x_id = x.id
+    elif isinstance(x, (Variable, Module)):
+      x_id = x._id
     else:
-      x_id = id(x)
+      return fn(x)
     if x_id not in refs:
       refs[x_id] = fn(x)
-    else:
-      pass
     return refs[x_id]
   return wrapped_fn
 
@@ -124,25 +125,25 @@ def get_arg_scope(x):
     nonlocal scopes
     if isinstance(x, Variable) and isinstance(x.scope, Scope):
       scopes.append(x.scope)
-      return VariablePlaceholder(x.collection, x.name, id(x))
+      return VariablePlaceholder(x.collection, x.name, x._id)
     elif isinstance(x, Module) and isinstance(x.scope, Scope):
-      x._try_setup(shallow=True)  # pylint: disable=protected-access
+      x._try_setup(shallow=True)
       scopes.append(x.scope)
       attrs = {
           f.name: getattr(x, f.name)
           for f in dataclasses.fields(x)
           if f.name != 'parent' and f.init
       }
       attrs = jax.tree_util.tree_map(get_arg_scope, attrs)
-      return InstancePlaceholder(x.__class__, attrs, id(x))
+      return InstancePlaceholder(x.__class__, attrs, x._id)
     return x
   new_args, new_kwargs = jax.tree_util.tree_map(get_arg_scope, (args, kwargs))
 
   # Gather scopes in Variables and Submodules passed as Module attributes.
   @functools.partial(_memoize_by_id, refs=refs)
   def get_scopes(module):
     nonlocal scopes
-    module._try_setup(shallow=True)  # pylint: disable=protected-access
+    module._try_setup(shallow=True)
     def get_scopes_inner(x):
       nonlocal scopes
       if isinstance(x, Module) and isinstance(x.scope, Scope):
@@ -303,9 +304,9 @@ 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', state.export())  # pylint: disable=protected-access
+        object.__setattr__(cloned, '_state', state.export())
         res = fn(cloned, *args, **kwargs)
-        self._state.reimport(cloned._state)  # pylint: disable=protected-access
+        self._state.reimport(cloned._state)
         _test_transformed_return_values(res, fn_name)
         return res
       # here we apply the given lifting transform to the scope-ingesting fn
@@ -351,9 +352,9 @@ 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', state.export())  # pylint: disable=protected-access
+      object.__setattr__(cloned, '_state', state.export())
       res = prewrapped_fn(cloned, *args, **kwargs)
-      self._state.reimport(cloned._state)  # pylint: disable=protected-access
+      self._state.reimport(cloned._state)
       _test_transformed_return_values(res, getattr(class_fn, '__name__', None))
       return res
     core_fns = [functools.partial(core_fn, prewrapped_fn, class_fn)
@@ -1325,8 +1326,8 @@ def wrapped_fn(self, *args, **kwargs):
     prewrapped_fn = wrap_method_once(class_fn)
     @functools.wraps(prewrapped_fn)
     def wrapped_fn(self, *args, **kwargs):
-      if ((not force and not linen_module._use_named_call)  # pylint: disable=protected-access
-          or self._state.in_setup):  # pylint: disable=protected-access
+      if ((not force and not linen_module._use_named_call)
+          or self._state.in_setup):
         return prewrapped_fn(self, *args, **kwargs)
       fn_name = class_fn.__name__
       method_suffix = f'.{fn_name}' if fn_name != '__call__' else ''
@@ -1335,9 +1336,9 @@ def wrapped_fn(self, *args, **kwargs):
       # make a scope-function to transform
       def core_fn(scopes, *args, **kwargs):
         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', self._state.export())
         res = prewrapped_fn(cloned, *args, **kwargs)
-        self._state.reimport(cloned._state)  # pylint: disable=protected-access
+        self._state.reimport(cloned._state)
         _test_transformed_return_values(res, fn_name)
         return res
       # here we apply the given lifting transform to the scope-ingesting fn

tests/linen/linen_test.py

@@ -14,8 +14,10 @@
 
 """Tests for flax.linen."""
 
+import copy
 from absl.testing import absltest, parameterized
 
+from flax import ids
 from flax import linen as nn
 
 import jax
@@ -361,5 +363,19 @@ def test_optimized_lstm_cell_matches_regular(self):
     jtu.check_eq(lstm_params, lstm_opt_params)
 
 
+class IdsTest(absltest.TestCase):
+
+  def test_hashable(self):
+    id1 = ids.uuid()
+    id2 = ids.uuid()
+    self.assertEqual(id1, id1)
+    self.assertNotEqual(id1, id2)
+    self.assertNotEqual(hash(id1), hash(id2))
+    id1c = copy.copy(id1)
+    id1dc = copy.deepcopy(id1)
+    self.assertNotEqual(hash(id1), hash(id1c))
+    self.assertNotEqual(hash(id1), hash(id1dc))
+
+
 if __name__ == '__main__':
   absltest.main()