pytreeclass (pypi): JAX compatible dataclass.
Keywords: python|machine-learning|pytorch|jax|dataclasses|deep-learning|tensorflow