jax_dataclasses
