Now running tree_check when initialising a Module

Author: patrick-kidger

equinox/_module.py

@@ -14,7 +14,7 @@
 from ._caches import internal_lru_caches
 from ._doc_utils import doc_repr
 from ._pretty_print import tree_pformat
-from ._tree import tree_equal
+from ._tree import tree_check_internal, tree_equal
 
 
 _P = ParamSpec("_P")
@@ -110,6 +110,11 @@ def _not_magic(k: str) -> bool:
 
 
 _has_dataclass_init = weakref.WeakKeyDictionary()
+_has_been_checked = weakref.WeakValueDictionary()
+
+
+def _skip(node):
+    return isinstance(node, Module) and node is _has_been_checked.get(id(node), None)
 
 
 # Inherits from ABCMeta as a convenience for a common use-case.
@@ -201,6 +206,10 @@ def __call__(cls, *args, **kwargs):
             else:
                 setattr(self, field.name, converter(getattr(self, field.name)))
         object.__setattr__(self, "__class__", cls)
+        # Note that this only runs during the initial creation, and not during
+        # unflattening.
+        tree_check_internal(self, _skip)
+        _has_been_checked[id(self)] = self
         return self
 
 

equinox/_tree.py

@@ -338,6 +338,14 @@ def is_leaf(node):
     return jtu.tree_flatten(pytree, is_leaf=is_leaf)
 
 
+def tree_check_internal(pytree, skip) -> None:
+    """As `tree_check`, but can skips checking some nodes (typically those that have
+    alread been checked).
+    """
+    all_nodes = {}
+    _tree_check(pytree, all_nodes, skip)
+
+
 def tree_check(pytree: Any) -> None:
     """Checks if the PyTree is well-formed: does it have no self-references, and does
     it have no duplicate layers.
@@ -389,13 +397,13 @@ def tree_check(pytree: Any) -> None:
     A `ValueError` if the PyTree is not well-formed.
     """
     all_nodes = {}
-    _tree_check(pytree, all_nodes)
+    _tree_check(pytree, all_nodes, skip=lambda _: False)
 
 
 _leaf_treedef = jtu.tree_structure(0)
 
 
-def _tree_check(node, all_nodes):
+def _tree_check(node, all_nodes, skip):
     subnodes, treedef = tree_flatten_one_level(node)
     # We allow duplicate leaves and empty containers, so don't raise an error with those
     if treedef != _leaf_treedef and treedef.num_leaves > 0:
@@ -424,5 +432,6 @@ def _tree_check(node, all_nodes):
             type_string = "<unknown type>"
         all_nodes[id(node)] = (True, type_string)
         for subnode in subnodes:
-            _tree_check(subnode, all_nodes)
+            if not skip(subnode):
+                _tree_check(subnode, all_nodes, skip)
         all_nodes[id(node)] = (False, type_string)

tests/test_module.py

@@ -1,4 +1,5 @@
 import functools as ft
+import gc
 from typing import Any
 
 import jax
@@ -252,3 +253,15 @@ def test_wrapped(getkey):
     y = eqx.filter_vmap(eqx.nn.Linear(2, 2, key=getkey()))
     x, y = eqx.filter((x, y), eqx.is_array)
     jtu.tree_map(lambda x, y: x + y, x, y)
+
+
+def test_tree_check_cache(getkey):
+    gc.collect()
+    has_been_checked = eqx._module._has_been_checked
+    num_checked = len(has_been_checked)
+    mlp = eqx.nn.MLP(2, 2, 2, 2, key=getkey())
+    # +4: one for `MLP`, and three for its `Linear` layers inside.
+    assert len(has_been_checked) == num_checked + 4
+    del mlp
+    gc.collect()
+    assert len(has_been_checked) == num_checked

tests/test_tree.py

@@ -231,9 +231,9 @@ def __init__(self, test=lambda x: 2 * x) -> None:
         def _transform(self, x):
             return self.test(x)
 
-    a = SubComponent()
     with pytest.raises(ValueError):
-        eqx.tree_check(a)
+        # Checking occurs when initialising the module.
+        SubComponent()
 
 
 def test_tree_check_none():