perf: dynamic (un)flattening code generation

[nstarman]

I haven't implemented the full logic, like the use of _FLATTEN_SENTINEL.

[nstarman]

@patrick-kidger This is ~80% of the way there. I think it's a tiny bit slower than the other PR but still gets equinox into the same ballpark as straight JAX.

The same test file, with the methods on the equinox module removed, should benchmark this PR.

This PR is missing the _flatten_sentinel and wrapper params logic. I'm AFK for most of the rest of the week, so if you like this approach and want to polish this off, that'd be awesome :).

(A third alternative would be to mypyc Module and let c eliminate the overhead.)

[patrick-kidger]

I really like this approach. If there's a need for greater speed then codegen sounds like a great way to get things to be fast-by-default.

I definitely don't have time to pursue this myself but I'll happily merge this once it's ready.

[nstarman]

For @patrick-kidger, a short report on current status:
I've cleaned up the PR code a lot and moved it to a separate module.
With the current PR https://github.com/user-attachments/files/22870041/overhead.ipynb.zip benchmarks great. It's about as fast as raw arrays and optimized pytrees. So that's good.
The part that I haven't been able to solve is the massive performance hit associated with getting and setting the WRAPPER_FIELD_NAMES. I've commented in all the changes necessary to use the wrapper fields.

[patrick-kidger]

Awesome! I really like the look of this PR.

The performance hit on the extra magic attributes – I think the common case is that none of these are present on the instance, and that instead they are class attributes. Perhaps this could be exploited, so that e.g. flattening check their presence and sets a boolean flag, and then unflattening skips this codepath if it is not required?

[nstarman]

The last commit needs a lot of cleanup, but if the '__module__' in obj.__dict__ guard is fine for filtering the wrapper attributes, then this PR gets us much closer to JAX speeds.
I'm clocking ~9.3 µs for the Equinox module versus ~7.7 µs for the JAX pytree in the benchmark notebook. It used to be 12-15 µs. Without the wrapper stuff it's ~8 µs, so there's still room for improvement.

[nstarman]

@patrick-kidger is '__module__' in obj.__dict__ sufficient a check?

[nstarman]

Ok. Final approximate timings on that performance notebook:

• Baseline is @partial(jax.tree_util.register_pytree_with_keys_class) ~ 7.7 µs
• Module before this PR ~ 12.5 µs
• Module after this PR ~ 10 µs
• Module after this PR, if I delete all wrapper field related code ~ 9 µs

So this PR is a ~50% improvement ((12.5 - 10) / (12.5 - 7.7)) but there's still another 2.3 µs to go.

[patrick-kidger]

Okay, many nits! But I basically really like this, it looks pretty much like exactly what we want to me.

[nstarman]

@patrick-kidger I think the only possibly unresolved comments are #1119 (comment) and #1119 (comment), but IMO they're good to go.

[nstarman]

I tested this on unxt and saw a nice µs speed boost in flattening/unflattening.
I plan to look at quax next for more performance enhancements.

[patrick-kidger]

Nice!
Given that this change occurs on the bottom of the tech stack, it would be good to also run the diffrax/optimistix/lineax tests against this too?

Just in case some of the weirder cases (overriding fields with properties etc) catch us out.

[nstarman]

uv add --upgrade  git+https://github.com/nstarman/equinox/@module-perf-2 
uv run pytest test -m "not slow" > diffrax_eqx1119_test.log 2>&1

diffrax_eqx1119_test.log

There is an error in diffrax, but when I re-run the tests outside of the custom equinox branch I get the same error, so I don't think it's related.

uv add --upgrade  git+https://github.com/nstarman/equinox/@module-perf-2 
uv run --extra tests pytest tests > ~optimistix_eqx1119_test.log 2>&1

optimistix_eqx1119_test.log

uv add --optional tests pytest beartype 
uv run --extra tests pytest tests > lineax_eqx1119_test.log 2>&1

lineax_eqx1119_test.log

[johannahaffner]

The diffrax errata were fixed in patrick-kidger/diffrax#696, and should not occur on dev, so I can confirm that these are unrelated :)

[patrick-kidger]

Awesome! Merged 🎉