Specify non-pytree node dataclass fields

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.