Repository: patrick-kidger/equinox
Pull request #994: perf: speed up module (un)flattening
Author: login: nstarman
Created: 2025-04-10T20:48:36Z | Merged: 2025-04-14T09:02:34Z

Partly addresses #961.

I don't see a benchmarking suite, e.g. using https://pypi.org/project/pytest-codspeed/, so in the attached notebook (attached as a ZIP because GH doesn't support directly attaching ipynb) I'm finding ~25% speedup towards matching the raw jax (around a 10% speedup overall).
[overhead.ipynb.zip](https://github.com/user-attachments/files/19694256/overhead.ipynb.zip)

This PR:

1. removes `_flatten_module` in favor of one defined in `_ActualModuleMeta.__new__` allowing for many aspects of the flattening to be pre-computed.
2. add slots to `_FlattenedData`, which should provide a minuscule speedup and reduction in memory use. Didn't hurt so 🤷.
3. Combined the `wrapper_field_names` and `wrapper_field_values` into a single field of `_FlattenedData` since it's only  ever used as a zipped combo. Also changed to pre-determined size for `wrapper_field_names_and_values` by storing the sentinel value. The sentinel is then used for skipping assignment when unflattening. In testing this did provide a speedup. 
4. Same as 3. for the static field names and values.


For future PRs...
- figure out how to not filter on `__dict__` so that `dynamic_fs = data_fs` and tuple sizes are pre-determined.
- try a namedtuple `_FlattenedData` to see if that speeds things up further.
- JAX dataclass-related speedups.



Diff:
diff --git a/equinox/_module.py b/equinox/_module.py
index 94ec55b3..69a1bc79 100644
--- a/equinox/_module.py
+++ b/equinox/_module.py
@@ -8,18 +8,19 @@
 import dataclasses
 import functools as ft
 import inspect
+import itertools
 import types
 import warnings
 import weakref
 from collections.abc import Callable
-from typing import Any, cast, Protocol, TYPE_CHECKING, TypeVar
+from typing import Any, cast, Final, Protocol, TYPE_CHECKING, TypeVar
 from typing_extensions import dataclass_transform, ParamSpec
 
 import jax
 import jax._src.traceback_util as traceback_util
 import jax.tree_util as jtu
 import numpy as np
-from jaxtyping import Array, Bool, PyTreeDef
+from jaxtyping import Array, Bool, PyTree, PyTreeDef
 
 from ._better_abstract import ABCMeta, dataclass
 from ._caches import cache_clears
@@ -133,6 +134,8 @@ class MyModule(eqx.Module):
 # This is the core of Equinox.
 #
 
+_flatten_sentinel: Final = doc_repr(object(), "-")
+
 
 @dataclass(frozen=True)
 class StrictConfig:
@@ -478,11 +481,69 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        sntl = _flatten_sentinel
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+
+            # Build dynamic keys, values in single pass
+            dynamic_fs = []
+            dynamic_vs = []
+            for k in data_fs:
+                v = d.get(k, sntl)
+                if v is sntl:
+                    continue
+                dynamic_fs.append(k)
+                dynamic_vs.append(v)
+
+            aux = _FlattenedData(
+                tuple(dynamic_fs),
+                tuple([(k, d.get(k, sntl)) for k in static_fs]),
+                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
+            )
+            return tuple(dynamic_vs), aux
+
+        def flatten_module_with_keys(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+
+            # Build dynamic keys, values in single pass
+            dynamic_fs = []
+            dynamic_vs = []
+            for k in data_fs:
+                v = d.get(k, sntl)
+                if v is sntl:
+                    continue
+                dynamic_fs.append(k)
+                dynamic_vs.append((jtu.GetAttrKey(k), v))
+
+            aux = _FlattenedData(
+                tuple(dynamic_fs),
+                tuple([(k, d.get(k, sntl)) for k in static_fs]),
+                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
+            )
+            return tuple(dynamic_vs), aux
+
         jtu.register_pytree_with_keys(
             cls,
-            flatten_with_keys=ft.partial(_flatten_module, with_keys=True),  # pyright: ignore
-            flatten_func=ft.partial(_flatten_module, with_keys=False),  # pyright: ignore
+            flatten_with_keys=flatten_module_with_keys,  # pyright: ignore
+            flatten_func=flatten_module,  # pyright: ignore
             unflatten_func=ft.partial(_unflatten_module, cls),  # pyright: ignore
         )
         # Done!
@@ -833,66 +894,14 @@ def bar(self):
 
 
 # Used to provide a pretty repr when doing `jtu.tree_structure(SomeModule(...))`.
-@dataclass()
+@dataclass(slots=True)
 class _FlattenedData:
     dynamic_field_names: tuple
-    static_field_names: tuple
-    static_field_values: tuple
-    wrapper_field_names: tuple
-    wrapper_field_values: tuple
+    static_fields: tuple[tuple[str, Any], ...]
+    wrapper_fields: tuple[tuple[str, Any], ...]
 
     def __repr__(self):
-        x = (
-            self.dynamic_field_names,
-            self.static_field_names,
-            self.static_field_values,
-        )
-        return repr(x)[1:-1]
-
-
-def _flatten_module(module: "Module", with_keys: bool):
-    # Subnodes in the PyTree
-    dynamic_field_names = []
-    dynamic_field_values = []
-    # Static metadata, placed in aux.
-    static_field_names = []
-    static_field_values = []
-    # Python metadata like `__doc__` and `__module__`.
-    wrapper_field_names = []
-    wrapper_field_values = []
-
-    for field_ in dataclasses.fields(module):
-        name = field_.name
-        try:
-            # Not `getattr` so that we don't pick up `property`s.
-            value = module.__dict__[name]
-        except KeyError:
-            # Uninitialised values during `__init__`, or when `property`s overwrite a
-            # field.
-            continue
-        if field_.metadata.get("static", False):
-            static_field_names.append(name)
-            static_field_values.append(value)
-        else:
-            dynamic_field_names.append(name)
-            if with_keys:
-                dynamic_field_values.append((jtu.GetAttrKey(name), value))
-            else:
-                dynamic_field_values.append(value)
-    sentinel = object()
-    for name in _wrapper_field_names:
-        value = module.__dict__.get(name, sentinel)
-        if value is not sentinel:
-            wrapper_field_names.append(name)
-            wrapper_field_values.append(value)
-    aux = _FlattenedData(
-        tuple(dynamic_field_names),
-        tuple(static_field_names),
-        tuple(static_field_values),
-        tuple(wrapper_field_names),
-        tuple(wrapper_field_values),
-    )
-    return tuple(dynamic_field_values), aux
+        return repr((self.dynamic_field_names, self.static_fields))[1:-1]
 
 
 def _unflatten_module(cls: type["Module"], aux: _FlattenedData, dynamic_field_values):
@@ -904,10 +913,9 @@ def _unflatten_module(cls: type["Module"], aux: _FlattenedData, dynamic_field_va
     module = object.__new__(cls)
     for name, value in zip(aux.dynamic_field_names, dynamic_field_values):
         object.__setattr__(module, name, value)
-    for name, value in zip(aux.static_field_names, aux.static_field_values):
-        object.__setattr__(module, name, value)
-    for name, value in zip(aux.wrapper_field_names, aux.wrapper_field_values):
-        object.__setattr__(module, name, value)
+    for name, value in itertools.chain(aux.static_fields, aux.wrapper_fields):
+        if value is not _flatten_sentinel:
+            object.__setattr__(module, name, value)
     return module
 
 
diff --git a/tests/test_pformat.py b/tests/test_pformat.py
index f8cc16c4..5597e88d 100644
--- a/tests/test_pformat.py
+++ b/tests/test_pformat.py
@@ -40,11 +40,11 @@ def g():
     i = jax.custom_vjp(f)
     j = jax.custom_jvp(f)
 
-    assert eqx.tree_pformat(f) == "<function f>"
-    assert eqx.tree_pformat(g) == "<wrapped function f>"
-    assert eqx.tree_pformat(h) == "partial(<function f>)"
-    assert eqx.tree_pformat(i) == "<function f>"
-    assert eqx.tree_pformat(j) == "<function f>"
+    assert eqx.tree_pformat(f) == "<function test_function.<locals>.f>"
+    assert eqx.tree_pformat(g) == "<wrapped function test_function.<locals>.f>"
+    assert eqx.tree_pformat(h) == "partial(<function test_function.<locals>.f>)"
+    assert eqx.tree_pformat(i) == "<function test_function.<locals>.f>"
+    assert eqx.tree_pformat(j) == "<function test_function.<locals>.f>"
 
 
 def test_struct_as_array():


Reviews:
user: login: patrick-kidger
body: This looks good to me! I've left some minor comments but I like the look of this.
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-11T16:32:42Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T18:37:54Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T18:38:04Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T20:24:33Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T20:25:05Z
user: login: patrick-kidger
body: Okay, this looks good to me! I think I spot one bug but otherwise I'm happy to merge this.
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-13T10:11:28Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-13T18:12:47Z
user: login: patrick-kidger
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-13T20:11:59Z
user: login: nstarman
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-13T21:45:21Z
user: login: patrick-kidger
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-14T08:42:33Z

Review comments:
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
path: equinox/_module.py
user: login: patrick-kidger
body: I think static fields could also fail to be initialised at times, e.g. when partway through `__init__`.

Also, I think this needs to be `obj.__dict__.get(name)` to avoid looking up through the class and the class hierarchy, in case of shadowing. We only want to grab instance attributes.
created_at: 2025-04-11T16:28:27Z
author_association: OWNER
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
+            # Subnodes in the PyTree (filtering on `__dict__` to avoid picking
+            # up `property` and uninitialized values).
+            dynamic_fs = [name for name in data_fs if name in obj.__dict__]
+            if with_keys:
+                dynamic_vs = [(jtu.GetAttrKey(k), getattr(obj, k)) for k in dynamic_fs]
+            else:
+                dynamic_vs = [getattr(obj, k) for k in dynamic_fs]
path: equinox/_module.py
user: login: patrick-kidger
body: Since speed is a concern, note that this is doing two lookups into `obj.__dict__`, once in `name in obj.__dict__` and another in `getattr`. I think the former could be skipped by replacing the latter with some kind of `obj.__dict__.get(k, _some_sentinel_value)`.

(I also speculate that a small speedup could be arranged by assigning to a local variable:
```python
get = obj.__dict__.get
dynamic_vs = [get(k, _some_kind_of_sentinel) for k in dynamic_fs]
```
)
created_at: 2025-04-11T16:30:40Z
author_association: OWNER
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
path: equinox/_module.py
user: login: nstarman
body: Done!
created_at: 2025-04-11T18:37:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
+            # Subnodes in the PyTree (filtering on `__dict__` to avoid picking
+            # up `property` and uninitialized values).
+            dynamic_fs = [name for name in data_fs if name in obj.__dict__]
+            if with_keys:
+                dynamic_vs = [(jtu.GetAttrKey(k), getattr(obj, k)) for k in dynamic_fs]
+            else:
+                dynamic_vs = [getattr(obj, k) for k in dynamic_fs]
path: equinox/_module.py
user: login: nstarman
body: Done!
created_at: 2025-04-11T18:38:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
path: equinox/_module.py
user: login: nstarman
body: https://github.com/patrick-kidger/equinox/pull/994/commits/53e829297c55c675cc49b3472a091219d141fc52 addresses this.
created_at: 2025-04-11T20:24:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +480,44 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /, with_keys: bool
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            # Static metadata, placed in aux.
+            static_vs = [getattr(obj, name) for name in static_fs]
+            # Subnodes in the PyTree (filtering on `__dict__` to avoid picking
+            # up `property` and uninitialized values).
+            dynamic_fs = [name for name in data_fs if name in obj.__dict__]
+            if with_keys:
+                dynamic_vs = [(jtu.GetAttrKey(k), getattr(obj, k)) for k in dynamic_fs]
+            else:
+                dynamic_vs = [getattr(obj, k) for k in dynamic_fs]
path: equinox/_module.py
user: login: nstarman
body: https://github.com/patrick-kidger/equinox/pull/994/commits/53e829297c55c675cc49b3472a091219d141fc52 only uses a single loop for getting from the dict.
created_at: 2025-04-11T20:25:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +480,77 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+            sntl = _flatten_sentinel
+
+            # Static fields, filtered on property / init state
+            static_vs = [v for k in static_fs if (v := d.get(k, sntl)) is not sntl]
path: equinox/_module.py
user: login: patrick-kidger
body: I think this will end up being of variable length if we get a sentinel value. That will then produce the wrong behaviour at unflattening time.

We could either similarly filter `static_fs` here, or skip pass sentinels in the unflattening. Wild guess that the latter is more efficient as in general sentinels will be uncommon.
created_at: 2025-04-13T10:09:31Z
author_association: OWNER
diff_hunk: @@ -478,11 +480,77 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+            sntl = _flatten_sentinel
+
+            # Static fields, filtered on property / init state
+            static_vs = [v for k in static_fs if (v := d.get(k, sntl)) is not sntl]
path: equinox/_module.py
user: login: nstarman
body: Done! I don't see any speedup, but %timeit is an imprecise tool.
This also afforded some further simplifications in `_FlattenedData` and then in `_unflatten_module`
created_at: 2025-04-13T18:12:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +481,69 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        sntl = _flatten_sentinel
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+
+            # Build dynamic keys, values in single pass
+            dynamic_fs = []
+            dynamic_vs = []
+            for k in data_fs:
+                v = d.get(k, sntl)
+                if v is sntl:
+                    continue
+                dynamic_fs.append(k)
+                dynamic_vs.append(v)
+
+            aux = _FlattenedData(
+                tuple(dynamic_fs),
+                tuple([(k, d.get(k, sntl)) for k in static_fs]),
+                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
path: equinox/_module.py
user: login: patrick-kidger
body: Nit: this is materialising a list and then converting it to a tuple. I think a direct `tuple((k, d.get(k, sntl) for k in static_fs)` should presumably save a small amount.
created_at: 2025-04-13T20:11:58Z
author_association: OWNER
diff_hunk: @@ -478,11 +481,69 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        sntl = _flatten_sentinel
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+
+            # Build dynamic keys, values in single pass
+            dynamic_fs = []
+            dynamic_vs = []
+            for k in data_fs:
+                v = d.get(k, sntl)
+                if v is sntl:
+                    continue
+                dynamic_fs.append(k)
+                dynamic_vs.append(v)
+
+            aux = _FlattenedData(
+                tuple(dynamic_fs),
+                tuple([(k, d.get(k, sntl)) for k in static_fs]),
+                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
path: equinox/_module.py
user: login: nstarman
body: One would think. But I've timed it out for different stuff and haven't found that to be the case. Here's a quick test:

![CleanShot 2025-04-13 at 17 44 27@2x](https://github.com/user-attachments/assets/0cad0000-bcbb-4549-967e-e0d8296914ed)

created_at: 2025-04-13T21:45:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -478,11 +481,69 @@ def __new__(
                                     "attempting to override "
                                     f"`{base.__module__}.{base.__qualname__}.{k}`."
                                 )
+
         # [Step 6] Register as a pytree.
+        data_fs = tuple(
+            f.name
+            for f in dataclasses.fields(cls)
+            if not f.metadata.get("static", False)
+        )
+        static_fs = tuple(
+            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
+        )
+
+        sntl = _flatten_sentinel
+
+        def flatten_module(
+            obj: "Module", /
+        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
+            """Flattens a module into a flattened representation."""
+            d = obj.__dict__
+
+            # Build dynamic keys, values in single pass
+            dynamic_fs = []
+            dynamic_vs = []
+            for k in data_fs:
+                v = d.get(k, sntl)
+                if v is sntl:
+                    continue
+                dynamic_fs.append(k)
+                dynamic_vs.append(v)
+
+            aux = _FlattenedData(
+                tuple(dynamic_fs),
+                tuple([(k, d.get(k, sntl)) for k in static_fs]),
+                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
path: equinox/_module.py
user: login: patrick-kidger
body: Mysterious! Well fair enough then :)
created_at: 2025-04-14T08:42:33Z
author_association: OWNER

PR comments:
user: login: nstarman
created_at: 2025-04-11T18:49:49Z
body: I refactored the module flattening for further optimizations.
~I think you need to enable the tests to run.~

There's still a lot of room for optimization! In the attached notebook JAX PyTrees take 10 us (vs 8.5 for raw arrays), and Equinox Module takes 12.5 us. 

Getting down to single loops might have improved performance at the cost of some verbosity. #998 helps with the verbosity, but is still a WIP to try to figure out some more performance-related gains.
author_association: CONTRIBUTOR
user: login: patrick-kidger
created_at: 2025-04-14T09:02:50Z
body: Alright, awesome -- merged! 😁 
author_association: OWNER