API Reference of lox

Core API

lox.log(data[, tags])	Fundamental logging primitive for Lox.
lox.tap(fun[, callback, argnames, tags, prefix])	A function transformation that taps into the execution of a JAX function and prints the values of specified arguments.
lox.spool(fun[, argnames, tags, keep_logs, ...])	Spools a function to extract logs generated during its execution.
lox.strip(fun[, argnames, tags])	Strips all logging operations from the given function by manipulating its Jaxpr.
lox.logdict(data, **steps)	A dictionary that stores log values and the steps at which they were recorded.
lox.save(data, path[, mode, key])	Save data to a specified path using a callback function.
lox.load(path[, key, result_shape_dtypes, ...])	Load data from a specified path.

Modules

lox.wandb
lox.loggers