lczerolens#

Main module for the lczerolens package.

Submodules#

Classes#

`LczeroBoard`	A class for wrapping the LczeroBoard class.
`InputEncoding`	Input encoding for the board tensor.
`LczeroModel`	Class for wrapping the LCZero models.
`Flow`	Base class for isolating a flow.
`Lens`	Generic lens class for analysing model activations.

Package Contents#

class lczerolens.LczeroBoard(fen=STARTING_FEN, *, chess960=False)[source]#

Bases: chess.Board

A class for wrapping the LczeroBoard class.

Parameters:

fen (Optional[str])
chess960 (bool)

static get_plane_order(us)[source]#

Get the plane order for the given us view.

Parameters:: us (bool) – The us_them tuple.
Returns:: The plane order.
Return type:: str

static get_piece_index(piece, us, plane_order=None)[source]#

Converts a piece to its index in the plane order.

Parameters:

piece (str) – The piece to convert.
us (bool) – The us_them tuple.
plane_order (Optional[str]) – The plane order.

Returns:

The index of the piece in the plane order.

Return type:

int

to_config_tensor(us=None, input_encoding=InputEncoding.INPUT_CLASSICAL_112_PLANE)[source]#

Converts a LczeroBoard to a tensor based on the pieces configuration.

Parameters:

us (Optional[bool]) – The us_them tuple.
input_encoding (InputEncoding) – The input encoding method.

Returns:

The 13x8x8 tensor.

Return type:

torch.Tensor

to_input_tensor(with_history=True, input_encoding=InputEncoding.INPUT_CLASSICAL_112_PLANE)[source]#

Create the lc0 input tensor from the history of a game.

Parameters:

with_history (bool) – Whether to include the history of the game.
input_encoding (InputEncoding) – The input encoding method.

Returns:

The 112x8x8 tensor.

Return type:

torch.Tensor

static encode_move(move, us)[source]#

Converts a chess.Move object to an index.

Parameters:

move (chess.Move) – The chess move to encode.
us (bool) – The side to move (True for white, False for black).

Returns:

The encoded move index.

Return type:

int

decode_move(index)[source]#

Converts an index to a chess.Move object.

Parameters:: index (int) – The index to convert.
Returns:: The chess move.
Return type:: chess.Move

get_legal_indices()[source]#

Gets the legal indices.

Returns:: Tensor containing indices of legal moves.
Return type:: torch.Tensor

get_next_legal_boards(n_history=7)[source]#

Gets the next legal boards.

Parameters:: n_history (int, optional) – Number of previous positions to keep in the move stack, by default 7.
Returns:: Generator yielding board positions after each legal move.
Return type:: Generator[LczeroBoard, None, None]

render_heatmap(heatmap, square=None, vmin=None, vmax=None, arrows=None, normalise='none', save_to=None, cmap_name='RdYlBu_r', alpha=1.0, flip_mode='board')[source]#

Render a heatmap on the board.

Parameters:

heatmap (torch.Tensor or numpy.ndarray) – The heatmap values to visualize on the board (64,) or (8, 8).
square (Optional[str], default=None) – Chess square to highlight (e.g. ‘e4’).
vmin (Optional[float], default=None) – Minimum value for the colormap normalization.
vmax (Optional[float], default=None) – Maximum value for the colormap normalization.
arrows (Optional[List[Tuple[str, str]]], default=None) – List of arrow tuples (from_square, to_square) to draw on board.
normalise (str, default="none") – Normalization method. Use “abs” for absolute value normalization.
save_to (Optional[str], default=None) – Path to save the visualization. If None, returns the figure.
cmap_name (str, default="RdYlBu_r") – Name of matplotlib colormap to use.
alpha (float, default=1.0) – Opacity of the heatmap overlay.
flip_mode (str, default="board") – Flip mode for black’s perspective. Use “board” to flip the board, “heatmap” to flip the heatmap.

Returns:

If save_to is None, returns (SVG string, matplotlib figure). If save_to is provided, saves files and returns None.

Return type:

Union[Tuple[str, matplotlib.figure.Figure], None]

Raises:

ValueError – If save_to is provided and does not end with .svg.

class lczerolens.InputEncoding[source]#

Bases: int, enum.Enum

Input encoding for the board tensor.

INPUT_CLASSICAL_112_PLANE = 0#

INPUT_CLASSICAL_112_PLANE_REPEATED = 1#

INPUT_CLASSICAL_112_PLANE_NO_HISTORY = 2#

class lczerolens.LczeroModel(model_key, *args, dispatch=False, meta_buffers=True, **kwargs)[source]#

Bases: nnsight.NNsight

Class for wrapping the LCZero models.

Parameters:

model_key (Union[str, torch.nn.Module])
dispatch (bool)
meta_buffers (bool)

trace(*inputs, **kwargs)[source]#

Entrypoint into the tracing and interleaving functionality nnsight provides.

In short, allows access to the future inputs and outputs of modules in order to trace what operations you would like to perform on them. This can be as simple as accessing and saving activations for inspection, or as complicated as transforming the activations and gradients in a forward pass over multiple inputs.

Parameters:

inputs (tuple[Any])
trace (bool, optional) – If to open a tracing context. Otherwise immediately run the model and return the raw output. Defaults to True.
invoker_args (Dict[str, Any], optional) – Keyword arguments to pass to Invoker initialization, and then downstream to the model’s .prepare_inputs(…) method. Used when giving input directly to .trace(…). Defaults to None.
kwargs (Dict[str, Any]) – Keyword arguments passed to Tracer initialization, and then downstream to the model’s ._execute(…) method.
backend (Union[Backend, str]) – Backend for this Tracer object.
remote (bool) – Use RemoteBackend with default url.

Raises:

ValueError – If trace is False and no inputs were provided (nothing to run with)

Returns:

Either the Tracer used for tracing, or the raw output if trace is False.

Return type:

Union[Tracer, Any]

Examples

There are a few ways you can use .trace(...) depending in your use case.

Lets use this extremely basic model for our examples:

import torch
from collections import OrderedDict

input_size = 5
hidden_dims = 10
output_size = 2

model = nn.Sequential(OrderedDict([
    ('layer1', torch.nn.Linear(input_size, hidden_dims)),
    ('sigma1', torch.nn.Sigmoid()),
    ('layer2', torch.nn.Linear(hidden_dims, output_size)),
    ('sigma2', torch.nn.Sigmoid()),
]))

example_input = torch.rand((1, input_size))

The first example has us running the model with a single example input, and saving the input and output of ‘layer2’ as well as the final output using the tracing context.

from nnsight import NNsight

with NNsight(model).trace(example_input) as model:

    l2_input = model.layer2.input.save()
    l2_output = model.layer2.output.save()

    output = model.output.save()

print(l2_input)
print(l2_output)
print(output)

The second example allows us to divide up multiple inputs into one batch, and scope an inner invoker context to each one. We indicate this simply by not passing and positional inputs into .trace(…). The Tracer object then expects you to enter each input via Tracer.invoke(…)

example_input2 = torch.rand((1, input_size))

with NNsight(model).trace() as model:

    with model.invoke(example_input):

        output1 = model.output.save()

    with model.invoke(example_input2):

        output2 = model.output.save()

print(output1)
print(output2)

_execute(*prepared_inputs, **kwargs)[source]#

Virtual method to run the underlying ._model with some inputs.

Default implementation util.applies moving all tensors to the device of the first parameter in ._model and passes the values into the model.

Parameters:: prepared_inputs (tuple[Any]) – Prepared inputs.
Return type:: Any

_prepare_inputs(*inputs, **kwargs)[source]#

Virtual method to prepare inputs before batching and execution and return batch size of prepared_inputs.

Default implementation just returns inputs and length of first input.

Parameters:

inputs (tuple[Any]) – Inputs to prepare for batching and execution.
int – Batch size of prepared_inputs.

Returns:

Prepared inputs, batch size of inputs.

Return type:

Tuple[tuple[Any], int]

__call__(*inputs, **kwargs)[source]#

__getattr__(key)[source]#

Wrapper of ._envoy’s attributes to access module’s inputs and outputs.

Returns:: Attribute.
Return type:: Any

__setattr__(key, value)[source]#: Overload setattr to create and set an Envoy when trying to set a torch Module.

property device#: Returns the device.

classmethod from_path(model_path)[source]#

Creates a wrapper from a model path.

Parameters:: model_path (str) – Path to the model file (.onnx or .pt)
Returns:: The wrapped model instance
Return type:: LczeroModel
Raises:: NotImplementedError – If the model file extension is not supported

classmethod from_onnx_path(onnx_model_path, check=True)[source]#

Builds a model from an ONNX file path.

Parameters:

onnx_model_path (str) – Path to the ONNX model file
check (bool, optional) – Whether to perform shape inference check, by default True

Returns:

The wrapped model instance

Return type:

LczeroModel

Raises:

FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded

classmethod from_torch_path(torch_model_path)[source]#

Builds a model from a PyTorch file path.

Parameters:

torch_model_path (str) – Path to the PyTorch model file

Returns:

The wrapped model instance

Return type:

LczeroModel

Raises:

FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded or is not a valid model type

_ensure_proper_forward()[source]#

class lczerolens.Flow(model_key, *args, **kwargs)[source]#

Bases: LczeroModel

Base class for isolating a flow.

_flow_type: str#

_registry: Dict[str, Type[Flow]]#

classmethod register(name)[source]#

Registers the flow.

Parameters:: name (str) – The name of the flow to register.
Returns:: Decorator function that registers the flow subclass.
Return type:: Callable
Raises:: ValueError – If the flow name is already registered.

classmethod from_name(name, *args, **kwargs)[source]#

Returns the flow from its name.

Parameters:

name (str) – The name of the flow to instantiate.
*args – Positional arguments passed to flow constructor.
**kwargs – Keyword arguments passed to flow constructor.

Returns:

The instantiated flow.

Return type:

Flow

Raises:

KeyError – If the flow name is not found.

classmethod from_model(name, model, *args, **kwargs)[source]#

Returns the flow from a model.

Parameters:

name (str) – The name of the flow to instantiate.
model (LczeroModel) – The model to create the flow from.
*args – Positional arguments passed to flow constructor.
**kwargs – Keyword arguments passed to flow constructor.

Returns:

The instantiated flow.

Return type:

Flow

Raises:

KeyError – If the flow name is not found.

classmethod is_compatible(model)[source]#

Checks if the model is compatible with this flow.

Parameters:: model (nn.Module) – The model to check compatibility with.
Returns:: Whether the model is compatible with this flow.
Return type:: bool

_ensure_proper_forward()[source]#: Rewrites the forward function to return the flow output.

class lczerolens.Lens(pattern=None)[source]#

Bases: abc.ABC

Generic lens class for analysing model activations.

Parameters:: pattern (Optional[str])

_lens_type: str#

_registry: Dict[str, Type[Lens]]#

classmethod register(name)[source]#

Registers the lens.

Parameters:: name (str) – The name of the lens.
Returns:: The decorator to register the lens.
Return type:: Callable
Raises:: ValueError – If the lens name is already registered.

classmethod from_name(name, *args, **kwargs)[source]#

Returns the lens from its name.

Parameters:: name (str) – The name of the lens.
Returns:: The lens instance.
Return type:: Lens
Raises:: KeyError – If the lens name is not found.

_pattern = None#

_reg_exp#

property pattern: str#

The pattern to match the modules.

Return type:: str

_get_modules(model)[source]#

Get the modules to intervene on.

Parameters:: model (lczerolens.model.LczeroModel)
Return type:: Generator[tuple[str, Any], None, None]

is_compatible(model)[source]#

Returns whether the lens is compatible with the model.

Parameters:: model (LczeroModel) – The NNsight model.
Returns:: Whether the lens is compatible with the model.
Return type:: bool

_ensure_compatible(model)[source]#

Ensure the lens is compatible with the model.

Parameters:: model (LczeroModel) – The NNsight model.
Raises:: ValueError – If the lens is not compatible with the model.

prepare(model, **kwargs)[source]#

Prepare the model for the lens.

Parameters:: model (LczeroModel) – The NNsight model.
Returns:: The prepared model.
Return type:: LczeroModel

abstract _intervene(model, **kwargs)[source]#

Intervene on the model.

Parameters:: model (LczeroModel) – The NNsight model.
Returns:: The intervention results.
Return type:: dict

_trace(model, *inputs, model_kwargs, intervention_kwargs)[source]#

Trace the model and intervene on it.

Parameters:

model (LczeroModel) – The NNsight model.
inputs (Union[LczeroBoard, torch.Tensor]) – The inputs.
model_kwargs (dict) – The model kwargs.
intervention_kwargs (dict) – The intervention kwargs.

Returns:

The intervention results.

Return type:

dict

analyse(model, *inputs, **kwargs)[source]#

Analyse the input.

Parameters:

model (LczeroModel) – The NNsight model.
inputs (Union[LczeroBoard, torch.Tensor]) – The inputs.

Returns:

The analysis results.

Return type:

dict

Raises:

ValueError – If the lens is not compatible with the model.

analyse_batched(model, iter_inputs, **kwargs)[source]#

Analyse a batches of inputs.

Parameters:

model (LczeroModel) – The NNsight model.
iter_inputs (Iterable[Tuple[Union[LczeroBoard, torch.Tensor], dict]]) – The iterator over the inputs.

Returns:

The iterator over the statistics.

Return type:

Generator[dict, None, None]

Raises:

ValueError – If the lens is not compatible with the model.