perf: add fastpaths in _ModuleMeta.__call__

[nstarman]

Scenario	Baseline	C1: init-check fastpath	C2: callable-warning scan	C3: wrapper-field setattr	C4: unchecked_init	C5: has_dataclass_init merge	Speedup
3 init=True fields	5.15	4.31	3.69	3.70	3.69	3.42	1.51x
1 converter field	3.52	3.28	2.83	2.85	2.82	2.55	1.38x
1 field + __check_init__	3.13	2.91	2.40	2.44	2.45	2.27	1.38x
3-level MRO + __check_init__	3.23	2.98	2.53	2.54	2.53	2.32	1.39x
10 init=True fields	11.80	9.20	9.30	9.50	9.27	8.78	1.34x
static field + 2 normal	6.00	5.53	4.83	4.90	4.74	3.64	1.65x
module_update_wrapper	11.86	11.75	11.76	10.52	10.47	9.36	1.27x
BoundMethod init (direct)	8.43	7.82	7.81	7.88	3.19	2.81	3.00x
attr access (x1)	8.77	8.17	8.29	8.24	3.58	3.23	2.72x
attr access (x3)	26.39	24.13	24.84	24.99	10.74	9.64	2.74x

import json
import sys
import timeit
import types
from pathlib import Path
from typing import Any

import equinox as eqx
from equinox._module._prebuilt import BoundMethod


# ---------------------------------------------------------------------------
# Module definitions
# ---------------------------------------------------------------------------


class _Simple(eqx.Module):
    a: int
    b: float
    c: str


class _WithConverter(eqx.Module):
    x: int = eqx.field(converter=int)


class _WithCheckInit(eqx.Module):
    a: int

    def __check_init__(self):
        pass


class _DeepA(eqx.Module):
    a: int

    def __check_init__(self):
        pass


class _DeepB(_DeepA):
    def __check_init__(self):
        pass


class _DeepC(_DeepB):
    def __check_init__(self):
        pass


class _ManyFields(eqx.Module):
    f0: Any
    f1: Any
    f2: Any
    f3: Any
    f4: Any
    f5: Any
    f6: Any
    f7: Any
    f8: Any
    f9: Any


class _StaticField(eqx.Module):
    name: str = eqx.field(static=True)
    value: int


class _Wrapped(eqx.Module):
    fn: Any

    def __call__(self, *args, **kwargs):
        return self.fn(*args, **kwargs)

    @property
    def __wrapped__(self):
        return self.fn


def _make_wrapped():
    return eqx.module_update_wrapper(_Wrapped(lambda x: x))


class _Host(eqx.Module):
    x: float
    y: float

    def method(self):
        return self.x + self.y

    def method2(self):
        return self.x * self.y

    def method3(self):
        return self.x - self.y


_host = _Host(1.0, 2.0)
_func: types.FunctionType = _Host.method  # type: ignore[assignment]


# ---------------------------------------------------------------------------
# Benchmark runner
# ---------------------------------------------------------------------------

SCENARIOS: list[tuple[str, str, dict[str, Any]]] = [
    ("simple", "_Simple(1, 2.0, 'x')", {}),
    ("converter", "_WithConverter(42)", {}),
    ("check_init", "_WithCheckInit(7)", {}),
    ("deep_inherit", "_DeepC(3)", {}),
    (
        "many_fields",
        "_ManyFields(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)",
        {},
    ),
    ("static_field", "_StaticField('hello', 99)", {}),
    ("module_update_wrapper", "_make_wrapped()", {}),
    (
        "BoundMethod_init",
        "BoundMethod(_func, _host)",
        {"BoundMethod": BoundMethod, "_func": _func, "_host": _host},
    ),
    (
        "attr_access_single",
        "_host.method",
        {"_host": _host},
    ),
    (
        "attr_access_three",
        "_host.method; _host.method2; _host.method3",
        {"_host": _host},
    ),
]

NUMBER = 200_000
REPEAT = 5


def run_benchmarks() -> dict[str, float]:
    g = {
        "_Simple": _Simple,
        "_WithConverter": _WithConverter,
        "_WithCheckInit": _WithCheckInit,
        "_DeepC": _DeepC,
        "_ManyFields": _ManyFields,
        "_StaticField": _StaticField,
        "_make_wrapped": _make_wrapped,
        "_host": _host,
        "_func": _func,
        "BoundMethod": BoundMethod,
    }
    results: dict[str, float] = {}
    for name, stmt, extra in SCENARIOS:
        times = timeit.repeat(
            stmt, globals={**g, **extra}, number=NUMBER, repeat=REPEAT
        )
        # µs per call
        results[name] = min(times) / NUMBER * 1e6
    return results


if __name__ == "__main__":
    out_path = Path(sys.argv[1]) if len(sys.argv) > 1 else None

    print(f"equinox version: {eqx.__version__}")
    print(f"number={NUMBER:,}  repeat={REPEAT}  (reporting min-of-repeats / call)\n")

    results = run_benchmarks()

    for name, us in results.items():
        print(f"  {name:<20s} {us:.4f} µs/call")

    if out_path is not None:
        out_path.write_text(json.dumps(results, indent=2))
        print(f"\nResults saved to {out_path}")

*AI disclosure: vibe engineered this. Looked over the code closely, but tests are Red/Green TDD.

[nstarman]

Ok. I think I've messed around with various ideas enough and this is PR is ready.
30% to 300% improvements for many common operations!

[patrick-kidger]

(Gosh, this took me a long time to get around to.)

So I don't love the unchecked_init flag.

1. This is going to be a maintenance headache: having two distinct codepaths is never a fun idea.
2. I strongly prefer to keep the eqx.Module implementation as readable as possible. It's already pretty complicated, and I much prefer it if a user can come in and understand what it does in each case.

I get the sense this might be a consistent difference between us! A lot of your (excellent!) PRs have focused on speed, but I am definitely happy to trade a little speed for a simplified maintenance story.

In particular in this case: initialising Equinox modules should happen pretty rarely in hot loops. Mostly this happens during model initialisation, or inside JIT'd regions. (And when crossing JIT boundaries, then we hit the flatten/unflattening API instead... which is an example of the opposite trade-off; here speed really does matter so much that I was happy to go as far as our codegen'd functions to order to optimize that case.)

WDYT?

[nstarman]

What about for the performance-sensitive classes like BoundMethod I just subclass _ModuleMeta. Then there's only one code path. The class flag was so that downstream classes can target enable faster construction, but getting BoundMethod is a higher priority for me.

_ModuleMeta(type)
Module(metaclass=_ModuleMeta)


_FastModuleMeta(_ModuleMeta)
BoundMethod(Module, metaclass=_FastModuleMeta)