traveler.trees

Classes

AttrPyTree(*args, **kwargs)	A JaxTree that allows for named attributes.
JaxDataset(*args, **kwargs)	A collection of JAX arrays of with metadata.
JaxTable(*args, **kwargs)	Analog of a pyarrow Table, but with JAX arrays.
WithTrace(main[, trace])	A class to hold auxiliary data in addition to the main tree.