diff --git a/equinox/_module/_flatten.py b/equinox/_module/_flatten.py index e116c19e..0a3dcb13 100644 --- a/equinox/_module/_flatten.py +++ b/equinox/_module/_flatten.py @@ -25,6 +25,7 @@ _INDENT: Final = " " * 4 +@jtu.register_static class _Missing: def __bool__(self): return False @@ -34,7 +35,6 @@ def __repr__(self): MISSING = _Missing() -jtu.register_pytree_node(_Missing, lambda _: ((), None), lambda _, __: MISSING) # Code template for flattening the wrapper fields