Chex dataclass defaulting mappable_dataclass=True

[balancap]

To start with, thanks for open sourcing your work on Chex, it's a great tooling library for building robust Jax applications!

As I was upgrading to the latest release 0.0.3, I noticed quite a few of my tests breaking. It happens that the default option mappable_dataclass=True in chex.dataclass is breaking the usual interface of dataclasses (which is clearly expected reading the code documentation!)

I guess probably from the perspective of Deepmind usage, it makes sense to default this option. But from an external user point of view, it is rather surprising to have a dataclass decorator not behaving like a dataclass. I think it would be great to make it clear in the library readme that this option needs to be turned off to get the full dataclass behaviour (or turned it off by default).

[hbq1]

Thanks for the positive feedback, we're happy that you found Chex useful in your projects!

The goal of Chex dataclasses is to extend the original interface rather than change it.
At this moment, the only existing incompatibility is that chex implementation prohibits positional arguments in constructor, which is done intentionally:
a) to reduce a risk of hard-to-detect errors related to wrong arguments ordering and
b) to improve code readability.
Apologies that this is not mentioned in the docs. Is it a reason of the breakage that you observed in your tests?

[balancap]

Thanks for the explanation. I completely agree with the motivation behind removing positional arguments. I guess I just got unlucky with using positional args on some very simple dataclasses + redefining the __getitem__ on some other for slicing purposes, which conflicted with the dict interface chex is setting up.

But no big deal, I guess mentioning it more clearly in the docs would avoid the surprise for others :)