Skip to content

WIP: faster low level modules#1221

Closed
nstarman wants to merge 3 commits into
patrick-kidger:mainfrom
nstarman:faster-low-level-modules
Closed

WIP: faster low level modules#1221
nstarman wants to merge 3 commits into
patrick-kidger:mainfrom
nstarman:faster-low-level-modules

Conversation

@nstarman

@nstarman nstarman commented Apr 21, 2026

Copy link
Copy Markdown
Contributor

Proof of concept.
Results in 21x speedup in BoundMethod initialization, translating into 14x speedup of method access from a Module.
This is stacked on top of speedups associated with #1220.

CleanShot 2026-04-21 at 12 16 36

In quax I was doing some line profiling and found that some examples in the tests had 2k method accesses, making this a non-trivial slowdown when considering a whole code pipeline.


This PR introduces a lower-level meta class and base class which internal Modules like BoundMethod can use, that bypasses the Module.__call__, and Module.__getattribte__. This is a WIP since there are details to be worked out, like:

  1. if instead it would be better to add a flag on Module that skips these slow methods, consolidating back to one class.
  2. if _AbstractModule should implement a from-style __setattr__, etc.

But I do think something should be done, since creating a BoundMethod on every method access is a common operation and is relatively slow, evidenced by this 14x performance improvement.

The main issue is the mildly(?) user-facing change

class MyModule(eqx.Module):
    a: int

    def f(self, b):
        return self.a + b

m = MyModule(13)
assert isinstance(m.f, eqx.Module)  # fails
assert isinstance(m.f, eqx.AbstractModule)  # if `_AbstractModule` is made public as `AbstractModule`

nstarman and others added 2 commits April 20, 2026 19:18
…e instantiation

_ModuleMeta.__call__ previously recomputed three things on every instantiation:
- dataclasses.fields(cls) — allocates a new tuple each call
- f.metadata.get("converter") / f.metadata.get("static") — per-field dict lookups
- for parent_cls in cls.__mro__: __check_init__ — full MRO scan each call

All three are now computed once at class-definition time in _ModuleMeta.__new__
and stored in four new _ModuleInfo fields:
- converter_fields: tuple of (name, converter) pairs
- static_field_names: tuple of static=True field names
- non_init_field_names: tuple of init=False field names
- check_init_methods: tuple of unbound __check_init__ functions from the MRO

__call__ now does a single _module_info[cls] lookup and iterates over
pre-filtered tuples, which are empty for the common case (no converters,
no static fields, no __check_init__).

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
@nstarman nstarman force-pushed the faster-low-level-modules branch from 1a5811b to de7e58b Compare April 21, 2026 17:04
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
@nstarman nstarman force-pushed the faster-low-level-modules branch from de7e58b to e29bd99 Compare April 21, 2026 17:26
@nstarman

Copy link
Copy Markdown
Contributor Author

@patrick-kidger what do you think of this potential fast-path for low-level Modules?

@patrick-kidger

Copy link
Copy Markdown
Owner

I think a __call__ fastpath would make sense, but rather than making it a separate class then I think we could arrange to have it be enabled automatically: for example, we don't need to check that fields are initialised if the class uses the dataclass-generated default; we can skip all the checks if all of the tuples are in the common case of being empty.

Ultimately I definitely don't want a new class – as you note it's user-visible, but more than that it's simply past the complexity/speed tradeoff for would consider valuable. It's important to me that Equinox be easy to maintain.

@nstarman

Copy link
Copy Markdown
Contributor Author

See #1230

@nstarman nstarman closed this May 14, 2026
@nstarman nstarman deleted the faster-low-level-modules branch May 14, 2026 04:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants