tree-math (conda): Mathematical operations for JAX pytrees
Keywords: jax|tree-math