etree.jax

[[Source]]

Expose 4 variants of the API depending on which backend is used.

jax = tree_utils.TreeAPI(backend.Jax())  # jax.tree_utils