Analog to flax.struct.dataclass

[homerjed]

Hi chex team,

Is there any potential for something like flax.struct.dataclass in chex?

Basically a kind of dataclass that can mark static arguments.

Two other variations include jax_dataclasses and simple-pytree though a chex official one would be cool for all its benefits of being part of chex.

Thank you!

[patrick-kidger]

Check out equinox.Module. This does dataclass + pytree registration as with the other cases, and supports static fields (eqx.static_field).

It also correctly handles a number of other edge cases: subclassing, docstring handling, bound methods are also pytrees, etc.

[homerjed]

Thanks @patrick-kidger, I will definitely use equinox next time around, its design is incredible.