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