lox.log

lox.log(data, tags=(), **steps)

Fundamental logging primitive for Lox. This primitive creates a logdict for a single data point and associates it with the provided steps.

Parameters:

• data (dict[str, Any]) – A dictionary containing the data to be logged.

• tags (Iterable[str]) – An iterable of strings representing tags associated with the log.

• steps (int) – Keyword arguments where keys are step names and values are step numbers.

Returns:

A logdict object containing the logged data and steps.

Return type:

logdict

Examples

By default lox.log complies with the pure functional programming paradigm of JAX, meaning it does not have side effects and does not mutate the state. Let’s look at this simple example of a function that adds 1.0 to its input.

>>> def f(x):
>>>     return x + 1.0
>>> f(1.0)
2.0

If we insert a logging statement inside the function, the function’s behavior remains unchanged.

>>> def f_log(x):
>>>     lox.log({"x": x})
>>>     return x + 1.0
>>> f_log(1.0)
2.0

In order to actually retrieve or use the logged data, we need to apply a function transformation such as spool or tap.