Public API: wax package

Subpackages

• wax.modules package
• wax.gym package
• wax.datasets package
• wax.universal package

wax.accessors	Define accessors for xarray and pandas data containers.
wax.compile	Compilation helper for Haiku Transformed and TransformedWithState pairs of pure funcions.
wax.encode	Encoding schemes to encode/decode numpy data types non supported by JAX, e.g.
wax.format	Format nested data structures to numpy/xarray/pandas containers.
wax.stream	Define Stream object used to synchronize in-memory data streams and unroll data transformations on it.
wax.transform	Transformation functions to work on batches of data.
wax.unroll	Unroll modules on data along first axis.
wax.utils	Some utils functions used in WAX-ML.