chex.dataclass vs eqx.Module vs linen.Module

[allen-adastra]

There are a lot of ways to solve the problem of registering classes as PyTrees. Torax currently uses chex.dataclass. I've been using eqx.Module in my own personal work for the most part. There's also flax.linen.Module as another option. I've noticed subtle differences when using the two; nothing huge. I'm wondering if anyone has insights into any fundamental pros/cons of the different approaches?

[araju]

chex.dataclass is lightweight and maybe the most common one used within our Google codebases. We mainly chose it for those reasons: it's simple, straightforward, and we are familiar with the gotchas.

eqx.Module is more powerful. I'll be honest, I actually haven't used it, so I cannot fully comment on the differences. One powerful feature we noticed right off the bat is the ability to partition and mark things as static, which is very useful for TORAX. We simply haven't gone all in on it yet because we need to test how it affects other parts of our TORAX ecosystem.

We didn't use flax.linen.Module simply because we haven't used it too much and it's inside flax, which we didn't want to pull in unnecessarily.

We see the choice here as an implementation detail, and we actually are going to be reevaluating this choice as we continue to expand TORAX.

Does this answer your question?

[allen-adastra]

Cool, thanks for the info!

It seems eqx.partition works for general PyTrees, so does not require eqx.Module in particular.

I've leaned against flax.linen.Module as well since it brings with it some additional requirements and deps, while chex.dataclass and eqx.Module just registers things as pytrees.

Had some discussion today with @dboyer-cfs about eqx.Module vs. chex.dataclass. We'll march forward with chex.dataclass for consistencies sake for now.