perf: enable module flatten/unflatten fastpath

[nstarman]

This PR enables power users to write their own flattening/unflattening procedures.
The advantage is that Module can be sped up to be as fast as custom pytree structures and jax arrays.
I'm attaching a Jupyter Notebook showing that in this example we achieve a ~50% speedup, which is ~100% of the overhead, making raw arrays, simple pytrees, and Module all equivalently fast.

I'm happy to add documentation / tests.

overhead.ipynb.zip

[patrick-kidger]

So Equinox looks to avoid special-casing any module methods. In particular it may be the case that someone already has a method called tree_unflatten etc, for some purpose unrelated to the tree map'ing of the module itself.

(The style you have here is actually what we used to have back in the distant early days of Equinox, and moved away from it.)

I'm not super what a better alternative is. Perhaps JAX might allow re-registering a PyTree with different flatten/unflatten rules.

[nstarman]

What about special-casing equinox-prefixed versions of these methods? eqx_tree_unflatten, etc. (they could even be private _eqx_tree_unflatten). That would be easy to support for fast-paths.

[nstarman]

@patrick-kidger if I can't crack #1119 (or even if I can, IMO it'd be nice to be able to customize) would _eqx_tree_unflatten be fine to add to Module?

[patrick-kidger]

So I'm really leaning against adding something like that to eqx.Module. Part of the design thesis of eqx.Module, as compared to jax.tree_util, is that custom flatten/unflatten functions are error-prone and simply aren't necessary – it suffices to just set dataclass fields instead.

Supposing #1119 comes good, what would be your use-case?

[nstarman]

E.g. not have the _MISSING (/flatten_sentinel) logic if I'm sure modules will be fully initialized. Avoid the wrapper stuff if needed. Write a mypyc transpiled mixin class with the (un)flattening logic that removes most (not all, because this would need @mypyc_attr(allow_interpreted_subclasses=True)) of the python overhead. For frequently-in-hot-loop classes like unxt.Quantity it would be nice to be able to achieve JAX speeds.

[nstarman]

And with #1119, it would be cool to show the user what the (un)flattening code is doing by attaching the generated functions to the classes!

class ModuleMeta:
    def __new__(...):
        cls._eqx_tree_flatten, cls._eqx_tree_unflatten = generate_functions(cls)
        jax.tree_util.register_stuff(cls._eqx_tree_flatten, cls._eqx_tree_unflatten)

class MyClass(eqx.Module):
    attr1: float

MyClass._eqx_tree_flatten?
>>> def flatten(self):
...  return (self.attr1,), ()

[nstarman]

I've rebased this PR on #1119.
And I've changed it so that the (un)flattening methods start with _eqx_

New approximate timings on that performance notebook:

• Baseline is @partial(jax.tree_util.register_pytree_with_keys_class) ~ 7.7 µs
• Module before perf: dynamic (un)flattening code generation #1119 ~ 12.5 µs
• Module after perf: dynamic (un)flattening code generation #1119 ~ 10 µs
• Module after perf: dynamic (un)flattening code generation #1119, if I delete all wrapper field related code ~ 9 µs
• This PR ~ 8.3 µs using a handwritten (un)flattening w/out the wrapper fields. So its ~1 µs faster and it's the power-user's choice for not including the wrapper field stuff :).

This PR allows for an ~88% improvement ((12.5 - 8.3) / (12.5 - 7.7)) if a power user wants to write these methods. The default (from #1119) is a ~50% improvement.

[nstarman]

@patrick-kidger, the before I rebase this PR and fix conflicts, I wanted to touch base r.e. the general idea.

The idea is that users can specify _eqx_tree_unflatten, _eqx_tree_flatten_with_keys, and _eqx_tree_flatten for
custom (un)flattening operations, and speedups of ~30%. Module falls back to the dynamically built methods if these are not provided.

        # Allow for classes to define their own (un)flattening procedures.
        unflatten = getattr(cls, "_eqx_tree_unflatten", False)
        flatten_with_k = getattr(cls, "_eqx_tree_flatten_with_keys", False)
        flatten = getattr(cls, "_eqx_tree_flatten", None)
        # Using `None` since it's the default value for
        # register_pytree_with_keys and has a truthy value of False.

        if unflatten and flatten_with_k:
            jtu.register_pytree_with_keys(
                cls,
                flatten_with_keys=flatten_with_k,
                flatten_func=flatten,
                unflatten_func=unflatten,
            )
        elif unflatten and flatten:
            jtu.register_pytree_node(cls, flatten, unflatten)
        else:
            # Partial or missing definition: fall back to automatic registration.
            if unflatten or flatten_with_k or flatten:
                warnings.warn(
                    "Class contains partial PyTree definition. "
                    "Falling back to automatic registration.",
                    stacklevel=2,
                )

            # Generate optimized flatten/unflatten functions
            flatten_func, flatten_with_keys_func, unflatten_func = (
                _generate_flatten_functions(cls, fields)
            )

            jtu.register_pytree_with_keys(
                cls,
                flatten_with_keys=flatten_with_keys_func,
                flatten_func=flatten_func,
                unflatten_func=ft.partial(unflatten_func, cls),
            )

[patrick-kidger]

So I think still hold to the idea from here:

So I'm really leaning against adding something like that to eqx.Module. Part of the design thesis of eqx.Module, as compared to jax.tree_util, is that custom flatten/unflatten functions are error-prone and simply aren't necessary – it suffices to just set dataclass fields instead.

if a user is willing to create custom flatten/unflatten then they should probably be using direct pytree registration, instead of using a Module. One of the big reasons to use Module, after all, is the creation of flatten/unflatten functions.

Conversely, if a user is using Module, then it's on us to make it as fast and performant as possible :) (As your work elsewhere in Equinox has been doing so well!)

[nstarman]

One of the big reasons to use Module, after all, is the creation of flatten/unflatten functions.

I just wish there were a good way to unregister things in JAX. Looking under the hood at the registration functions, it's such a complex, variable set of dicts that I imagine it's hard to do. So offering alt registrations here is much easier.
Is there any custom path you'd feel comfortable with? A class-kwarg to identify which are the (un)flattening functions?

[patrick-kidger]

I think the correct custom-(un)flattening path is to use JAX's custom pytree-registration directly, and not to use eqx.Module at all. Equinox Modules are basically just pytrees+dataclasses, if the user wants custom pytrees then Equinox isn't actually giving you a whole lot beyond @dataclasses.dataclass.

[nstarman]

Methods-as-pytrees, the downstream ecosystem that does isinstance(x, Module), .... there's a lot of ways in which a Chex object differs.
Given quax's heavy use of Module at lower-levels of the JAX plumbing than most other packages, quax is one of the more performance-sensitive equinox use cases. Hence these scattershot PRs ratcheting up performance.
The three most consequential areas are 1. initialization, 2. flattening, and 3. accessing methods.
Flattening has gotten a lot faster, but if Module supported some kind of customization we could achieve equal performance to other libraries and hand-coded dataclasses.

The 2 ideas I have are:

1. This PR for special methods on class.
2. New class arguments for providing these functions, e.g.

class MyClass(Module, unflatten=..., flatten=...): ...

I'm very happy to close this PR in favour of the other idea, if that'd be better.