treeo (pypi): A small library for creating and manipulating custom JAX Pytree classes
Keywords: jax|pytree