Specify non-pytree node dataclass fields

[ethanluoyc]

Hi,

Thanks for making this awesome library!

Is it possible to specify fields in the chex.dataclass definitions to not include certain fields? This is a feature supported in flax https://flax.readthedocs.io/en/latest/_modules/flax/struct.html#dataclass
which I found to be quite useful when defining data classes with fields (such as JAX functions) that shouldn't be mapped over with dm-tree or jax.tree_map. I am not sure if this is supported out of the box by chex at the moment but is something that I hope would be part of chex.

[felixchalumeau]

Hi guys,
As this issue ever been discussed? Any recent update?
Thanks!

[pablo2909]

Would love to see this feature brought to Chex :)

[felixchalumeau]

Hi guys 👋

Any update on this? Is there any plan to have it in the future?