lczerolens.lens

Generic lens class.

class lczerolens.lens.Lens(pattern: str | None = None)

Bases: ABC

Generic lens class for analysing model activations.

analyse(model: LczeroModel, *inputs: LczeroBoard | Tensor, **kwargs) → dict

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: LczeroModel, iter_inputs: Iterable[LczeroBoard | Tensor], **kwargs) → Generator[dict, None, None]

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.

classmethod from_name(name: str, *args, **kwargs) → Lens

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.

is_compatible(model: LczeroModel) → bool

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

property pattern: str

The pattern to match the modules.

prepare(model: LczeroModel, **kwargs) → LczeroModel

Prepare the model for the lens.

Parameters:

model (LczeroModel) – The NNsight model.

Returns:

The prepared model.

Return type:

LczeroModel

classmethod register(name: str) → Callable

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.