GitHub repository: stefansommer/hyperiax
Description: Tree traversals using JAX