treex (pypi): A Pytree Module system for Deep Learning in JAX
Keywords: deep-learning|jax