Load models and encode inputs

LczeroModel owns neural execution while LczeroEvaluator owns chess-aware preparation and output standardization. This notebook creates and reloads a tiny local PyTorch network so every cell runs offline. The same evaluator boundary accepts converted ONNX files and Hub models.

from pathlib import Path
from tempfile import TemporaryDirectory

import chess
import torch
from torch import nn

from lczerolens import InputFormat, LczeroEvaluator, LczeroModel


class TinyNetwork(nn.Module):
    def forward(self, planes):
        batch = planes.shape[0]
        policy = torch.zeros((batch, 1858), device=planes.device)
        wdl = torch.tensor([0.4, 0.3, 0.3], device=planes.device).expand(batch, -1).clone()
        mlh = torch.full((batch,), 12.0, device=planes.device)
        return policy, wdl, mlh


with TemporaryDirectory() as directory:
    model_path = Path(directory) / "tiny-network.pt"
    torch.save(TinyNetwork(), model_path)
    model = LczeroModel.from_path(str(model_path), out_keys=["policy", "wdl", "mlh"])

model.heads, model.network is not None, model.network_checksum[:12]

(('policy', 'wdl', 'mlh'), True, 'sha256:f861c')

For a real network, use LczeroModel.from_path("network.onnx") or install the hub extra and call LczeroModel.from_hf("organization/model"). Pin the Hub revision when provenance must identify immutable weights.

boards = [chess.Board(), chess.Board()]
boards[1].push_uci("e2e4")
evaluator = LczeroEvaluator(model, input_format=InputFormat.CLASSICAL_112)
prepared = evaluator.prepare(boards)
{
    "batch_size": tuple(prepared.batch_size),
    "planes_shape": tuple(prepared["input", "planes"].shape),
    "legal_moves": prepared["input", "legal_mask"].sum(-1).tolist(),
    "device": str(prepared.device),
}

{'batch_size': (2,),
 'planes_shape': (2, 112, 8, 8),
 'legal_moves': [20, 20],
 'device': 'cpu'}

The prepared TensorDict is the advanced execution boundary. External instrumentation can add nested keys between prepare() and finish(); lczerolens validates its own keys without discarding those additions. Device movement belongs to the model and TensorDict.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
evaluator.model.to(device)
prepared = evaluator.prepare(boards)
prepared["demo", "position_id"] = torch.arange(len(boards), device=device).unsqueeze(-1)
evaluations = evaluator.finish(boards, evaluator.model(prepared))
assert ("demo", "position_id") in evaluations.tensors.keys(include_nested=True, leaves_only=True)
{
    "heads": evaluator.model.heads,
    "value_origin": evaluations[0].value.origin.value,
    "mlh": evaluations[0].mlh,
    "retained_demo_key": evaluations.tensors["demo", "position_id"].tolist(),
}

{'heads': ('policy', 'wdl', 'mlh'),
 'value_origin': 'derived_from_wdl',
 'mlh': 12.0,
 'retained_demo_key': [[0], [1]]}