Analog to `flax.struct.dataclass`

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`](https://github.com/brentyi/jax_dataclasses) and [`simple-pytree`](https://github.com/cgarciae/simple-pytree/blob/main/simple_pytree/pytree.py) though a `chex` official one would be cool for all its benefits of being part of `chex`.

Thank you!