traveler.trees

Generic JAX pytree containers and tracing helpers.

Classes

AttrPyTree(*args, **kwargs)	A JaxTree that allows for named attributes.
WithTrace(main[, trace])	A class to hold auxiliary data in addition to the main tree.