Modules

Clouseau tracks the forward pass of a model and saves the intermediate arrays to a file.

Currently it supports Jax and Pytorch models. They way this is achieved is different for both frameworks.

• For Jax it uses a wrapper class that wraps each (callable) node in the pytree ( see also https://github.com/patrick-kidger/equinox/issues/864). Saving arrays to file is a side effect in Jax. See e.g. https://docs.jax.dev/en/latest/external-callbacks.html However the global cache seems acceptable in combination with jax.experimental.io_callback, which is explicitly designed for this purpose.
• For Pytorch it uses a forward hook that is registered on each module. See e.g. https://web.stanford.edu/~nanbhas/blog/forward-hooks-pytorch/ The forward hook is de-registered after the forward pass is done.

In both cases Jax / Pytorch it tracks the output of the layer. I might add tracking of the inputs as well later.

is_jax_model(model)

Check if model is a jax model.

Any Python object is a valid pytree (a bare object flattens to itself as a single leaf), so jax.tree.flatten alone would accept almost anything, e.g. an int. We additionally require model to be a non-leaf pytree node that carries at least one jax.Array leaf, i.e. a container with parameters that clouseau can actually record.

Source code in clouseau/inspector.py

def is_jax_model(model: AnyModel) -> bool:
    """Check if model is a jax model.

    Any Python object is a valid pytree (a bare object flattens to itself as a
    single leaf), so ``jax.tree.flatten`` alone would accept almost anything,
    e.g. an ``int``. We additionally require ``model`` to be a non-leaf pytree
    node that carries at least one ``jax.Array`` leaf, i.e. a container with
    parameters that clouseau can actually record.
    """
    try:
        import jax
    except ImportError:
        return False

    leaves, treedef = jax.tree.flatten(model)

    # a bare leaf (int, str, plain object, or a raw array) flattens to itself
    if treedef.num_leaves == 1 and leaves[0] is model:
        return False

    return any(isinstance(leaf, jax.Array) for leaf in leaves)

is_torch_model(model)

Check if model is a torch model

Source code in clouseau/inspector.py

def is_torch_model(model: AnyModel) -> bool:
    """Check if model is a torch model"""
    try:
        import torch

        return isinstance(model, torch.nn.Module)
    except ImportError:
        return False

magnify(filename=DEFAULT_PATH / 'activations-000.safetensors', framework='numpy', device=None)

Visualize nested arrays using treescope

Source code in clouseau/inspector.py

def magnify(
    filename: str | Path = DEFAULT_PATH / "activations-000.safetensors",
    framework: str = "numpy",
    device: Any = None,
) -> None:
    """Visualize nested arrays using treescope"""
    data = read_from_safetensors(filename, framework=framework, device=device)

    print_tree(data)

tail(model, path=DEFAULT_PATH, filter_=None, is_leaf=None, filename_pattern='activations-{idx:03d}.safetensors', max_size_mb=1024)

Tail and record the forward pass of a model

Parameters:

Name	Type	Description	Default
model	object	The model to inspect. Can be a PyTorch model or JAX/Equinox model.	required
path	str or Path	Path where to store the forward pass arrays.	DEFAULT_PATH
filter_	callable	Function that filters which tensors to inspect. Takes the pytree leaves, child modules as input and returns a boolean.	None
is_leaf	callable	Function that determines whether a node in the model tree should be treated as a leaf. Takes a node as input and returns a boolean. If True, the node will not be traversed further. This is particularly useful for JAX/Equinox models to control the granularity of inspection.	None

Returns:

Type	Description
_Inspector	Inspector instance that can be used as a context manager.

Examples:

>>> import torch
>>> from clouseau import inspector, magnifier
>>> model = torch.nn.Linear(10, 5)
>>> with inspector.tail(model,  path=".clouseau/trace-torch.safetensors") as fmodel:
...     out = fmodel(torch.randn(3, 10))

When working with a JAX/Equinox model, it is important to add .block_until_ready()

>>> import jax
>>> with inspector.tail(model, path=".clouseau/trace-jax.safetensors") as fmodel:
...     fmodel(x, time).block_until_ready()

Source code in clouseau/inspector.py

def tail(
    model: AnyModel,
    path: str | Path = DEFAULT_PATH,
    filter_: Callable[[Any, Any], bool] | None = None,
    is_leaf: Callable[[Any, Any], bool] | None = None,
    filename_pattern: str = "activations-{idx:03d}.safetensors",
    max_size_mb: int = 1024,
) -> _Recorder:
    """Tail and record the forward pass of a model

    Parameters
    ----------
    model : object
        The model to inspect. Can be a PyTorch model or JAX/Equinox model.
    path : str or Path
        Path where to store the forward pass arrays.
    filter_ : callable
        Function that filters which tensors to inspect.
        Takes the pytree leaves, child modules as input and returns a boolean.
    is_leaf : callable, optional
        Function that determines whether a node in the model tree should be treated as a leaf.
        Takes a node as input and returns a boolean. If True, the node will not be traversed further.
        This is particularly useful for JAX/Equinox models to control the granularity of inspection.

    Returns
    -------
    _Inspector
        Inspector instance that can be used as a context manager.

    Examples
    --------
    >>> import torch
    >>> from clouseau import inspector, magnifier
    >>> model = torch.nn.Linear(10, 5)
    >>> with inspector.tail(model,  path=".clouseau/trace-torch.safetensors") as fmodel:
    ...     out = fmodel(torch.randn(3, 10))

    When working with a JAX/Equinox model, it is important to add `.block_until_ready()`
    >>> import jax
    >>> with inspector.tail(model, path=".clouseau/trace-jax.safetensors") as fmodel:
    ...     fmodel(x, time).block_until_ready()

    """
    return _Recorder(
        model=model,
        path=path,
        filter_=filter_,
        is_leaf=is_leaf,
        max_size_mb=max_size_mb,
        filename_pattern=filename_pattern,
    )