diff --git a/docs/sphinx/examples/qec/python/tn_noise_learning.py b/docs/sphinx/examples/qec/python/tn_noise_learning.py index f1631b93c..71fa61f60 100644 --- a/docs/sphinx/examples/qec/python/tn_noise_learning.py +++ b/docs/sphinx/examples/qec/python/tn_noise_learning.py @@ -21,14 +21,12 @@ # [Begin Documentation] """ -Noise learning with NMOptimizer on a Stim repetition-code circuit. +Online noise learning with NMOptimizer from a Stim detector error model. -This script demonstrates how to use NMOptimizer to fit per-error noise -probabilities to syndrome data sampled from a Stim repetition-code memory -experiment. Starting from uniform initial priors, Adam optimization on -logits drives the cross-entropy loss down toward the true DEM error rates, -and a held-out evaluation compares the learned model's logical error rate -against the static uniform-prior baseline. +This example demonstrates how to fit per-error noise probabilities using +fresh syndrome batches sampled from a Stim detector error model. The optimizer +starts from uniform priors, resamples training data each iteration, and then +compares held-out logical error rate for uniform, learned, and true DEM priors. Requirements: pip install cudaq-qec[tensor-network-decoder] stim beliefmatching @@ -39,9 +37,16 @@ import stim from beliefmatching.belief_matching import detector_error_model_to_check_matrices -import cudaq_qec as qec from cudaq_qec import NMOptimizer, make_compiled_step +TRAIN_SHOTS = 5000 +EVAL_SHOTS = 20000 +ITERS = 100 +LR = 1e-2 +DTYPE = "float64" +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" +EXECUTE = "codegen" + def parse_detector_error_model(dem): matrices = detector_error_model_to_check_matrices(dem) @@ -49,14 +54,40 @@ def parse_detector_error_model(dem): matrices.check_matrix.astype(np.float64).toarray(out=H) L = np.zeros(matrices.observables_matrix.shape) matrices.observables_matrix.astype(np.float64).toarray(out=L) - priors = [float(p) for p in matrices.priors] + priors = np.array([float(p) for p in matrices.priors], dtype=np.float64) return H, L, priors +def make_sampler(circuit, seed): + try: + return circuit.compile_detector_sampler(seed=seed) + except TypeError: + return circuit.compile_detector_sampler() + + +def sample_batch(sampler, shots): + det_events, obs_flips = sampler.sample(shots, separate_observables=True) + return det_events.astype(float), obs_flips.ravel().astype(bool) + + +def probs_to_logits(probs): + probs = np.clip(probs, 1e-7, 1.0 - 1e-7) + return np.log(probs / (1.0 - probs)) + + +def evaluate_ler(H, L, priors, det_events, obs_flips): + opt = NMOptimizer(H, + L, + priors.tolist(), + det_events, + obs_flips, + dtype=DTYPE, + device=DEVICE, + execute=EXECUTE) + return opt.logical_error_rate() + + def main(): - # Asymmetric noise (data 10x measurement) so the uniform initial - # prior is meaningfully wrong and the optimizer has signal to - # learn; with symmetric noise, uniform is already near-optimal. circuit = stim.Circuit.generated( "repetition_code:memory", rounds=5, @@ -66,92 +97,68 @@ def main(): ) dem = circuit.detector_error_model(decompose_errors=True) H, L, true_priors = parse_detector_error_model(dem) - true_probs = np.array(true_priors) n_checks, n_errors = H.shape print(f"DEM: {n_checks} checks, {n_errors} errors") - print(f"True priors: mean={true_probs.mean():.4e} " - f"min={true_probs.min():.4e} max={true_probs.max():.4e} " - f"(spread {true_probs.max() / true_probs.min():.1f}x)") + print(f"True priors: mean={true_priors.mean():.4e} " + f"min={true_priors.min():.4e} max={true_priors.max():.4e}") + print(f"Device: {DEVICE}, execute={EXECUTE}") - num_shots = 1000 - sampler = circuit.compile_detector_sampler() - det_events, obs_flips = sampler.sample(num_shots, separate_observables=True) - det_events = det_events.astype(float) - obs_flips = obs_flips.ravel().astype(bool) + train_sampler = make_sampler(circuit, seed=1234) + det_events, obs_flips = sample_batch(train_sampler, TRAIN_SHOTS) - uniform = float(true_probs.mean()) + uniform = np.full(n_errors, true_priors.mean(), dtype=np.float64) opt = NMOptimizer(H, - L, [uniform] * n_errors, + L, + uniform.tolist(), det_events, obs_flips, - dtype="float64") - - # Optimize in logit space — numerically stabler than raw probs. - def _to_logits(p): - p = np.clip(p, 1e-7, 1 - 1e-7) - return -np.log(1.0 / p - 1.0) - - logits = torch.tensor( - _to_logits(np.full(n_errors, uniform)), - dtype=torch.float64, - device=opt.torch_device, - requires_grad=True, - ) - adam = torch.optim.Adam([logits], lr=1e-2) + dtype=DTYPE, + device=DEVICE, + execute=EXECUTE) + + logits = torch.tensor(probs_to_logits(uniform), + dtype=getattr(torch, DTYPE), + device=opt.torch_device, + requires_grad=True) + adam = torch.optim.Adam([logits], lr=LR) step_fn = make_compiled_step(opt, logits, adam) - iters = 300 - losses = [float(step_fn().detach().cpu()) for _ in range(iters)] + losses = [] + mae_history = [] + true_probs = torch.tensor(true_priors, + dtype=getattr(torch, DTYPE), + device=opt.torch_device) + for it in range(ITERS): + if it > 0: + det_events, obs_flips = sample_batch(train_sampler, TRAIN_SHOTS) + opt.update_dataset(det_events, obs_flips) + loss = step_fn() + learned_probs = torch.sigmoid(logits) + losses.append(float(loss.detach().cpu())) + mae_history.append( + float( + torch.mean(torch.abs(learned_probs - + true_probs)).detach().cpu())) + learned = torch.sigmoid(logits).detach().cpu().numpy() - print(f"Loss: {losses[0]:.2f} -> {losses[-1]:.2f} " - f"({iters} Adam steps)") - print(f"True priors: mean={true_probs.mean():.4e} " - f"min={true_probs.min():.4e} max={true_probs.max():.4e}") + print(f"Loss: {losses[0]:.4e} -> {losses[-1]:.4e} " + f"({ITERS} online Adam steps, {TRAIN_SHOTS} shots/step)") + print(f"Prior MAE: {mae_history[0]:.4e} -> {mae_history[-1]:.4e}") print(f"Learned priors: mean={learned.mean():.4e} " f"min={learned.min():.4e} max={learned.max():.4e}") - if losses[-1] >= losses[0]: - raise RuntimeError(f"Training did not reduce loss at all: " - f"{losses[0]:.2f} -> {losses[-1]:.2f}") - - # Held-out LER comparison is the real gate: a noise model is only - # useful if it decodes better than uniform priors. 20k shots keeps - # the per-run std of the (static - learned) difference around 0.001, - # so the +0.002 gate sits many sigmas below the expected gain even - # without a fixed RNG seed. - num_test = 20000 - test_events, test_flips = sampler.sample(num_test, - separate_observables=True) - test_events = test_events.astype(float) - test_flips_bool = test_flips.ravel().astype(bool) - - def _ler(noise: list[float]) -> float: - decoder = qec.get_decoder( - "tensor_network_decoder", - H, - logical_obs=L, - noise_model=noise, - contract_noise_model=True, - ) - res = decoder.decode_batch(test_events) - pred = np.array([r.result[0] > 0.5 for r in res], dtype=bool) - return float(np.mean(pred != test_flips_bool)) - - ler_static = _ler([uniform] * n_errors) - ler_learned = _ler(learned.tolist()) - - print(f"LER (static uniform priors): {ler_static:.4f} ({num_test} shots)") - print( - f"LER (learned priors): {ler_learned:.4f} ({num_test} shots)") - print(f"Absolute improvement: {ler_static - ler_learned:+.4f}") - - min_improvement = 0.002 - if ler_static - ler_learned < min_improvement: - raise RuntimeError( - f"Learned LER ({ler_learned:.4f}) did not beat the static " - f"baseline ({ler_static:.4f}) by at least {min_improvement:.4f}.") + eval_sampler = make_sampler(circuit, seed=4321) + eval_events, eval_flips = sample_batch(eval_sampler, EVAL_SHOTS) + ler_uniform = evaluate_ler(H, L, uniform, eval_events, eval_flips) + ler_learned = evaluate_ler(H, L, learned, eval_events, eval_flips) + ler_true = evaluate_ler(H, L, true_priors, eval_events, eval_flips) + + print(f"Held-out LER (uniform priors): {ler_uniform:.4f}") + print(f"Held-out LER (learned priors): {ler_learned:.4f}") + print(f"Held-out LER (true DEM priors): {ler_true:.4f}") + print(f"Absolute LER improvement: {ler_uniform - ler_learned:+.4f}") if __name__ == "__main__": diff --git a/libs/qec/python/cudaq_qec/__init__.py b/libs/qec/python/cudaq_qec/__init__.py index 310785e7d..f3206c922 100644 --- a/libs/qec/python/cudaq_qec/__init__.py +++ b/libs/qec/python/cudaq_qec/__init__.py @@ -156,7 +156,9 @@ def iter_namespace(ns_pkg): try: from .plugins.decoders.tensor_network_utils.nm_optimizer import ( NMOptimizer, + make_batch_sliced_step, make_compiled_step, + make_training_step, ) except (ModuleNotFoundError, ImportError): pass diff --git a/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/contractors.py b/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/contractors.py index 5ded4fb3a..3d7951d0c 100644 --- a/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/contractors.py +++ b/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/contractors.py @@ -7,11 +7,13 @@ # ============================================================================ # from __future__ import annotations +from collections import OrderedDict from collections.abc import Callable from dataclasses import dataclass, field from typing import Any, ClassVar import opt_einsum as oe +import torch from quimb.tensor import TensorNetwork @@ -33,6 +35,40 @@ def contractor(subscripts: str, return oe.contract(subscripts, *tensors, optimize=optimize) +def oe_torch_contractor(subscripts: str, + tensors: list[torch.Tensor], + optimize: str = "auto", + **_: Any) -> Any: + """Perform einsum contraction using opt_einsum with torch tensors. + + Execution follows the input tensor device, so CUDA tensors stay on + CUDA while still preserving torch autograd. + """ + return oe.contract(subscripts, *tensors, optimize=optimize, backend="torch") + + +_OE_EXPR_CACHE_MAXSIZE = 32 +_oe_expr_cache: OrderedDict[tuple, Any] = OrderedDict() + + +def oe_torch_compiled_contractor(subscripts: str, + tensors: list[torch.Tensor], + optimize: str = "auto", + **_: Any) -> Any: + """Perform einsum contraction with a cached opt_einsum expression.""" + shapes = tuple(t.shape for t in tensors) + key = (subscripts, shapes, str(optimize)) + if key in _oe_expr_cache: + _oe_expr_cache.move_to_end(key) + else: + if len(_oe_expr_cache) >= _OE_EXPR_CACHE_MAXSIZE: + _oe_expr_cache.popitem(last=False) + _oe_expr_cache[key] = oe.contract_expression(subscripts, + *shapes, + optimize=optimize) + return _oe_expr_cache[key](*tensors, backend="torch") + + def cutn_contractor(subscripts: str, tensors: list[Any], optimize: Any | None = None, @@ -109,6 +145,11 @@ class ContractorConfig: _allowed_configs: ClassVar[tuple[tuple[str, str, str], ...]] = ( ("numpy", "numpy", "cpu"), ("torch", "torch", "cpu"), + ("torch", "torch", "cuda"), + ("oe_torch", "torch", "cpu"), + ("oe_torch", "torch", "cuda"), + ("oe_torch_compiled", "torch", "cpu"), + ("oe_torch_compiled", "torch", "cuda"), ("cutensornet", "numpy", "cuda"), ("cutensornet", "torch", "cuda"), ) @@ -116,6 +157,8 @@ class ContractorConfig: _contractors: ClassVar[dict[str, Callable]] = { "numpy": contractor, "torch": contractor, + "oe_torch": oe_torch_contractor, + "oe_torch_compiled": oe_torch_compiled_contractor, "cutensornet": cutn_contractor, } diff --git a/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/nm_optimizer.py b/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/nm_optimizer.py index 4e908908f..6101b1f3d 100644 --- a/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/nm_optimizer.py +++ b/libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_utils/nm_optimizer.py @@ -18,8 +18,9 @@ from __future__ import annotations import warnings -from typing import Any, Literal +from typing import Any, Callable, Literal +import cudaq_qec as qec import numpy as np import numpy.typing as npt import opt_einsum as oe @@ -27,7 +28,9 @@ from quimb.tensor import TensorNetwork from ..tensor_network_decoder import TensorNetworkDecoder +from .noise_models import factorized_noise_model from .tensor_network_factory import ( + tensor_network_from_parity_check, tensor_network_from_syndrome_batch, prepare_syndrome_data_batch, ) @@ -41,6 +44,7 @@ "float32": 1e-6, } _SUPPORTED_DTYPES: tuple[str, ...] = ("float32", "float64") +_TORCH_EINSUM_MAX_DIMS = 25 def _validate_and_clamp_priors(noise_model: Any, dtype: str) -> list[float]: @@ -106,6 +110,12 @@ def _clamp_log_input(x: torch.Tensor) -> torch.Tensor: return x.clamp_min(torch.finfo(x.dtype).tiny) +def _normalize_prediction(x: torch.Tensor) -> torch.Tensor: + """Normalize decoder scores while keeping zero rows finite.""" + norm = x.sum(dim=1, keepdim=True).clamp_min(torch.finfo(x.dtype).tiny) + return x / norm + + def remap_eq_to_ascii(eq: str) -> str: """Rewrite an einsum equation so every label is in ``[a-zA-Z]``. @@ -147,6 +157,90 @@ def remap_eq_to_ascii(eq: str) -> str: return f"{out_lhs}->{''.join(out_rhs_chars)}" +def _path_largest_intermediate(info: Any) -> float: + value = getattr(info, "largest_intermediate", None) + if value is None: + return float("inf") + try: + return float(value) + except (TypeError, ValueError, OverflowError): + return float("inf") + + +def _path_opt_cost(info: Any) -> float: + value = getattr(info, "opt_cost", None) + if value is None: + return float("inf") + try: + return float(value) + except (TypeError, ValueError, OverflowError): + return float("inf") + + +def _einsum_rank(eq: str) -> int: + """Maximum tensor rank touched by one explicit einsum step.""" + if "->" not in eq: + return 0 + lhs, rhs = eq.split("->", 1) + terms = [term for term in lhs.split(",") if term] + terms.append(rhs) + return max((len(term) for term in terms), default=0) + + +def _path_max_einsum_rank(info: Any) -> int: + """Maximum rank emitted by an opt_einsum contraction path.""" + ranks = [ + _einsum_rank(step[2]) for step in getattr(info, "contraction_list", ()) + ] + return max(ranks, default=0) + + +def _select_default_torch_path( + eq: str, + shapes: tuple[tuple[int, ...], ...], +) -> tuple[Any, Any]: + candidates: list[tuple[str, Any, Any]] = [] + optimizers: list[tuple[str, Any]] = [ + ("greedy", "greedy"), + ("auto", "auto"), + ("auto-hq", "auto-hq"), + ("random-greedy", "random-greedy"), + ("random-greedy-128", "random-greedy-128"), + ] + + for tag, optimize in optimizers: + try: + path, info = oe.contract_path(eq, + *shapes, + shapes=True, + optimize=optimize) + except Exception as exc: + warnings.warn( + f"NMOptimizer path candidate {tag!r} failed: {exc!r}", + RuntimeWarning, + stacklevel=3, + ) + continue + candidates.append((tag, path, info)) + + if not candidates: + raise RuntimeError("No NMOptimizer contraction path candidate " + "succeeded.") + + safe_candidates = [ + candidate for candidate in candidates + if _path_max_einsum_rank(candidate[2]) <= _TORCH_EINSUM_MAX_DIMS + ] + if safe_candidates: + candidates = safe_candidates + + _tag, selected_path, selected_info = min( + candidates, + key=lambda c: (_path_largest_intermediate(c[2]), _path_opt_cost(c[2])), + ) + return selected_path, selected_info + + class NMOptimizer(TensorNetworkDecoder): """Differentiable noise-model optimiser for the TN decoder. @@ -189,7 +283,8 @@ class NMOptimizer(TensorNetworkDecoder): execute: Forward backend. ``"codegen"`` (default) partial-evaluates the path into a flat Python function; ``"unrolled"`` keeps an interpretive einsum list; ``"opt_einsum"`` dispatches via - :func:`opt_einsum.contract_expression`. + :func:`opt_einsum.contract_expression`; ``"cutensornet"`` + executes the reduced contraction with cuTensorNet on CUDA. compile_mode: Forwarded to :func:`torch.compile`; ignored when ``compile=False``. dynamic_syndromes: If ``True`` (default), syndromes are runtime @@ -199,6 +294,9 @@ class NMOptimizer(TensorNetworkDecoder): the closure as constants -- fewer runtime einsums, but every :meth:`update_dataset` call rebuilds the graph. Only affects ``execute="codegen"``. + Per-error noise tensors are contracted with their adjacent code + tensors using differentiable torch ops, then the reduced + network is contracted with the selected ``execute`` backend. Example (logit-space, no clamping needed):: @@ -227,11 +325,12 @@ def __init__( device: str = "cuda", *, compile: bool = False, - execute: Literal["codegen", "unrolled", "opt_einsum"] = "codegen", + execute: Literal["codegen", "unrolled", "opt_einsum", + "cutensornet"] = "codegen", compile_mode: str | None = None, dynamic_syndromes: bool = True, ) -> None: - if execute not in ("unrolled", "opt_einsum", "codegen"): + if execute not in ("unrolled", "opt_einsum", "codegen", "cutensornet"): raise ValueError(f"Invalid execute mode: {execute!r}") if dtype not in _SUPPORTED_DTYPES: raise ValueError(f"Invalid dtype {dtype!r}; expected one of " @@ -240,42 +339,50 @@ def __init__( # Sanitise once so the base TN tensors and ``self._noise_probs`` # see identical values (see :func:`_validate_and_clamp_priors`). noise_model = _validate_and_clamp_priors(noise_model, dtype) - - super().__init__( - H, - logical_obs, - noise_model, - check_inds=check_inds, - error_inds=error_inds, - logical_inds=logical_inds, - logical_tags=logical_tags, - contract_noise_model=False, - dtype=dtype, - device=device, - ) - - # Force the torch backend so tensor data lives in the autograd - # graph (the base class would otherwise pick cutensornet/numpy - # on GPU). Contractions still go through codegen / cotengra - # below, not cuTensorNet. - if self.contractor_config.contractor_name == "cutensornet" \ - and self.contractor_config.backend != "torch": + target_device = device + if "cuda" in target_device and not torch.cuda.is_available(): warnings.warn( - "NMOptimizer requires the torch backend for autograd; " - f"switching contractor backend from " - f"{self.contractor_config.backend!r} to 'torch'. " - "Contractions are executed via codegen/opt_einsum, not " - "cuTensorNet.", - stacklevel=3, - ) - self._set_contractor( - "cutensornet", - self.contractor_config.device, - "torch", - dtype, + "CUDA was requested for NMOptimizer, but torch CUDA is not " + "available. Using CPU.", + UserWarning, + stacklevel=2, ) + target_device = "cpu" + + if execute == "cutensornet" and "cuda" not in target_device: + raise ValueError("execute='cutensornet' requires a CUDA device.") + + # Build the topology directly so NMOptimizer never selects the + # TensorNetworkDecoder cuTensorNet contractor during setup. + qec.Decoder.__init__(self, H) + + num_checks, num_errs = H.shape + if check_inds is None: + self.check_inds = [f"s_{j}" for j in range(num_checks)] + else: + assert len(check_inds) == num_checks, ( + f"check_inds must have length {num_checks}, " + f"but got {len(check_inds)}.") + self.check_inds = check_inds + if error_inds is None: + self.error_inds = [f"e_{j}" for j in range(num_errs)] + else: + assert len(error_inds) == num_errs, ( + f"error_inds must have length {num_errs}, " + f"but got {len(error_inds)}.") + self.error_inds = error_inds + + self.logical_obs_inds = ["obs"] + self.parity_check_matrix = H.copy() + self.code_tn = tensor_network_from_parity_check( + self.parity_check_matrix, + col_inds=self.error_inds, + row_inds=self.check_inds, + ) + self.replace_logical_observable(logical_obs, + logical_inds=logical_inds, + logical_tags=logical_tags) - # Swap the base's placeholder single-syndrome TN for a batched one. self._syndrome_tags = [f"SYN_{i}" for i in range(len(self.check_inds))] self.syndrome_tn = tensor_network_from_syndrome_batch( syndrome_data, @@ -285,14 +392,23 @@ def __init__( ) self._batch_size = syndrome_data.shape[0] - # Re-stitch ``full_tn`` around the batched syndrome TN. self.full_tn = TensorNetwork() self.full_tn = self.full_tn.combine(self.code_tn, virtual=True) self.full_tn = self.full_tn.combine(self.logical_tn, virtual=True) self.full_tn = self.full_tn.combine(self.syndrome_tn, virtual=True) - self.full_tn = self.full_tn.combine(self.noise_model, virtual=True) - self._set_tensor_type(self.syndrome_tn) + self.path_single = "auto" + self.path_batch = "auto" + self.slicing_batch = tuple() + self.slicing_single = tuple() + + contractor = ("cutensornet" + if execute == "cutensornet" else "oe_torch_compiled") + self._set_contractor(contractor, target_device, "torch", dtype) + self.init_noise_model( + factorized_noise_model(self.error_inds, noise_model), + contract=False, + ) torch_dtype = getattr(torch, self._dtype) self._noise_probs = torch.tensor( @@ -301,10 +417,10 @@ def __init__( device=self.torch_device, requires_grad=True, ) - # The base's noise tensors stay in ``full_tn`` as numpy - # placeholders: ``_snapshot_arrays_and_eq`` uses ``id()`` to - # locate their positions, then ``self._noise_probs`` (autograd - # live) is written into those slots. Do not strip them. + # The base's noise tensors stay in ``full_tn`` as placeholders: + # ``_snapshot_arrays_and_eq`` uses ``id()`` to locate their + # positions, then ``self._noise_probs`` (autograd live) is written + # into those slots. Do not strip them. self._suspend_loss_rebuild = True self.observable_flips = observable_flips @@ -443,117 +559,307 @@ def _as_torch(x): self._syndrome_shapes: tuple[tuple[int, ...], ...] = tuple( tuple(s.shape) for s in self._syndrome_arrays) - if self._execute_mode == "opt_einsum": - shapes = tuple(t.shape for t in tensors) - self._oe_expr = oe.contract_expression( - self._eq_batch, - *shapes, - optimize=self.path_batch - if self.path_batch not in (None, "auto") else "auto", - ) - self._path_steps = None - else: - self._oe_expr = None - # Flatten the path into ``[(eq, idxs, sorted_desc), ...]``; - # ``sorted_desc`` is precomputed for the unrolled-mode pop - # walk, and labels are remapped to ASCII because - # opt_einsum falls back to unicode past 52 indices and - # torch.einsum rejects those. - shapes = tuple(t.shape for t in tensors) - _, info = oe.contract_path( - self._eq_batch, - *shapes, - shapes=True, - optimize=self.path_batch - if self.path_batch not in (None, "auto") else "auto", - ) - self._path_steps = [(remap_eq_to_ascii(step[2]), tuple(step[0]), - tuple(sorted(step[0], reverse=True))) - for step in info.contraction_list] + self._oe_expr = None + self._path_steps = None + self._build_reduced_tn_state() self._compile_predict() self._compile_loss() def _compile_predict(self) -> None: """Build ``self._predict_fn`` for the configured execute mode.""" - builders = { - "opt_einsum": self._build_predict_opt_einsum, - "unrolled": self._build_predict_unrolled, - "codegen": self._build_predict_codegen, - } - self._predict_fn = builders[self._execute_mode]() + self._predict_fn = self._build_predict_reduced() self._compiled_predict = self._maybe_torch_compile(self._predict_fn, kind="predict") - def _build_predict_opt_einsum(self): - """opt_einsum-backed predict: reuse the cached contract expression.""" - static_arrays = self._static_arrays - syndrome_positions = tuple(p for p, _t in self._syndrome_positions) - noise_pos_ordered = self._noise_pos_ordered - n = len(self._tensors_ref) - oe_expr = self._oe_expr + def _build_reduced_tn_state(self) -> None: + """Build reduced TN topology plus differentiable noise recipes.""" + from collections import defaultdict + + error_inds_set = set(self.error_inds) + survivor_lookup: dict[tuple[tuple[str, ...], frozenset[str]], int] = {} + doomed_lookup: dict[tuple[tuple[str, ...], frozenset[str]], int] = {} + for opt_pos, tensor in enumerate(self._tensors_ref): + key = (tuple(tensor.inds), frozenset(tensor.tags)) + if any(ind in error_inds_set for ind in tensor.inds): + doomed_lookup[key] = opt_pos + else: + survivor_lookup[key] = opt_pos + + reduced_tn = self.full_tn.copy() + recipes: list[dict[str, Any]] = [] + merged_id_to_recipe_idx: dict[int, int] = {} + + for error_idx, error_ind in enumerate(self.error_inds): + doomed = [t for t in reduced_tn.tensors if error_ind in t.inds] + code_tensors = [t for t in doomed if "NOISE" not in t.tags] + code_opt_positions = [ + doomed_lookup[(tuple(t.inds), frozenset(t.tags))] + for t in code_tensors + ] + + ids_before = {id(t) for t in reduced_tn.tensors} + reduced_tn.contract_ind(error_ind) + new_tensors = [ + t for t in reduced_tn.tensors if id(t) not in ids_before + ] + assert len(new_tensors) == 1 + new_tensor = new_tensors[0] + merged_id_to_recipe_idx[id(new_tensor)] = error_idx + + out_inds = tuple(new_tensor.inds) + mapping = {error_ind: "e"} + next_code = ord("a") + for ind in out_inds: + while chr(next_code) == "e": + next_code += 1 + mapping[ind] = chr(next_code) + next_code += 1 + + noise_str = mapping[error_ind] + code_strs = [ + "".join(mapping[ind] for ind in t.inds) for t in code_tensors + ] + out_str = "".join(mapping[ind] for ind in out_inds) + ordered_code_positions: list[int] = [None] * len( # type: ignore + code_tensors) + for tensor, opt_pos in zip(code_tensors, code_opt_positions): + non_error_ind = next( + ind for ind in tensor.inds if ind != error_ind) + ordered_code_positions[out_inds.index(non_error_ind)] = opt_pos + + recipes.append({ + "eq": ",".join([noise_str] + code_strs) + "->" + out_str, + "ordered_code_positions": ordered_code_positions, + "k": len(code_tensors), + }) + + reduced_eq = reduced_tn.get_equation( + output_inds=("batch_index", self.logical_obs_inds[0])) + reduced_shapes = tuple(t.shape for t in reduced_tn.tensors) + + reduced_static: dict[int, torch.Tensor] = {} + reduced_syndrome: list[tuple[int, int]] = [] + reduced_recipes: dict[int, int] = {} + syn_pos_to_idx = { + p: i for i, (p, _) in enumerate(self._syndrome_positions) + } + for pos, tensor in enumerate(reduced_tn.tensors): + if id(tensor) in merged_id_to_recipe_idx: + reduced_recipes[pos] = merged_id_to_recipe_idx[id(tensor)] + continue + + key = (tuple(tensor.inds), frozenset(tensor.tags)) + opt_pos = survivor_lookup[key] + if opt_pos in self._static_arrays: + reduced_static[pos] = self._static_arrays[opt_pos] + elif opt_pos in syn_pos_to_idx: + reduced_syndrome.append((pos, syn_pos_to_idx[opt_pos])) + else: + raise AssertionError( + f"Reduced tensor at position {pos} maps to full tensor " + f"position {opt_pos}, which is not static or syndrome.") + + reduced_path = self.path_batch + if reduced_path in (None, "auto"): + reduced_path, reduced_info = _select_default_torch_path( + reduced_eq, reduced_shapes) + else: + _path, reduced_info = oe.contract_path(reduced_eq, + *reduced_shapes, + shapes=True, + optimize=reduced_path) + reduced_path_steps = [(remap_eq_to_ascii(step[2]), tuple(step[0]), + tuple(sorted(step[0], reverse=True))) + for step in reduced_info.contraction_list] + reduced_oe_expr = None + if self._execute_mode == "opt_einsum": + reduced_oe_expr = oe.contract_expression(reduced_eq, + *reduced_shapes, + optimize=reduced_path) + + recipe_to_reduced_pos = {ri: pos for pos, ri in reduced_recipes.items()} + reduced_recipe_positions = tuple( + recipe_to_reduced_pos[ri] for ri in range(len(recipes))) + groups_by_k: dict[int, list[int]] = defaultdict(list) + for recipe_idx, recipe in enumerate(recipes): + groups_by_k[recipe["k"]].append(recipe_idx) + + batched_groups: list[dict[str, Any]] = [] + device = self.torch_device + for k, error_indices in sorted(groups_by_k.items()): + out_letters: list[str] = [] + next_code = ord("a") + for _ in range(k): + while chr(next_code) in ("e", "n"): + next_code += 1 + out_letters.append(chr(next_code)) + next_code += 1 + + if k == 0: + eq = "ne->ne" + else: + check_strs = [f"n{letter}e" for letter in out_letters] + eq = "ne," + ",".join(check_strs) + "->n" + "".join(out_letters) + + stacked_checks = [] + for axis in range(k): + axis_arrays = [ + self._static_arrays[recipes[ri]["ordered_code_positions"] + [axis]] for ri in error_indices + ] + stacked_checks.append(torch.stack(axis_arrays, dim=0)) + + batched_groups.append({ + "k": + k, + "eq": + eq, + "error_indices_t": + torch.tensor(error_indices, dtype=torch.long, + device=device), + "stacked_checks": + stacked_checks, + "recipe_indices": + error_indices, + }) + + self._batched_einsum_groups = batched_groups + self._reduced_eq = reduced_eq + self._reduced_static_positions = reduced_static + self._reduced_syndrome_positions = reduced_syndrome + self._reduced_recipe_positions = reduced_recipe_positions + self._reduced_oe_expr = reduced_oe_expr + self._reduced_path_steps = reduced_path_steps + self._reduced_n_tensors = len(reduced_tn.tensors) + self._reduced_n_recipes = len(recipes) + self._reduced_info = reduced_info + self.path_batch = reduced_path + self.slicing_batch = tuple() + + def _build_predict_reduced(self): + """Predict using the reduced TN plus batched noise precontraction.""" + builders = { + "opt_einsum": self._build_predict_reduced_opt_einsum, + "unrolled": self._build_predict_reduced_unrolled, + "codegen": self._build_predict_reduced_codegen, + "cutensornet": self._build_predict_reduced_cutensornet, + } + return builders[self._execute_mode]() + + def _precontract_reduced_noise( + self, noise_probs: torch.Tensor) -> tuple[torch.Tensor, ...]: + noise_stacked = torch.stack((1.0 - noise_probs, noise_probs), dim=-1) + recipe_arrays = [None] * self._reduced_n_recipes + for group in self._batched_einsum_groups: + noise_batch = noise_stacked[group["error_indices_t"]] + if group["k"] == 0: + out_batch = noise_batch + else: + out_batch = torch.einsum(group["eq"], noise_batch, + *group["stacked_checks"]) + for i, recipe_idx in enumerate(group["recipe_indices"]): + recipe_arrays[recipe_idx] = out_batch[i] + return tuple(recipe_arrays) # type: ignore[return-value] + + def _build_predict_reduced_opt_einsum(self): + static_positions = self._reduced_static_positions + syndrome_positions = self._reduced_syndrome_positions + oe_expr = self._reduced_oe_expr + recipe_positions = self._reduced_recipe_positions + n = self._reduced_n_tensors def _predict(noise_probs: torch.Tensor, syndrome_tuple: tuple[torch.Tensor, ...]) -> torch.Tensor: - noise_stacked = torch.stack((1.0 - noise_probs, noise_probs), - dim=-1) arrays: list[torch.Tensor] = [None] * n # type: ignore - for pos, arr in static_arrays.items(): + for pos, arr in static_positions.items(): arrays[pos] = arr - for pos, arr in zip(syndrome_positions, syndrome_tuple): + for pos, syndrome_idx in syndrome_positions: + arrays[pos] = syndrome_tuple[syndrome_idx] + for pos, arr in zip(recipe_positions, + self._precontract_reduced_noise(noise_probs)): arrays[pos] = arr - for k, pos in enumerate(noise_pos_ordered): - arrays[pos] = noise_stacked[k] - # Torch backend is auto-selected from the tensor type; - # avoids the per-call ``backend=`` dispatch. out = oe_expr(*arrays) - return out / out.sum(dim=1, keepdim=True) + return _normalize_prediction(out) return _predict - def _build_predict_unrolled(self): - """Unrolled predict: walk the cached pairwise contraction path.""" - static_arrays = self._static_arrays - syndrome_positions = tuple(p for p, _t in self._syndrome_positions) - noise_pos_ordered = self._noise_pos_ordered - n = len(self._tensors_ref) - path_steps = self._path_steps + def _build_predict_reduced_cutensornet(self): + static_positions = self._reduced_static_positions + syndrome_positions = self._reduced_syndrome_positions + recipe_positions = self._reduced_recipe_positions + eq = self._reduced_eq + n = self._reduced_n_tensors + contractor = self.contractor_config.contractor + device_id = self.contractor_config.device_id + + def _predict(noise_probs: torch.Tensor, + syndrome_tuple: tuple[torch.Tensor, ...]) -> torch.Tensor: + arrays: list[torch.Tensor] = [None] * n # type: ignore + for pos, arr in static_positions.items(): + arrays[pos] = arr + for pos, syndrome_idx in syndrome_positions: + arrays[pos] = syndrome_tuple[syndrome_idx] + for pos, arr in zip(recipe_positions, + self._precontract_reduced_noise(noise_probs)): + arrays[pos] = arr + out = contractor(eq, + arrays, + optimize=self.path_batch, + slicing=self.slicing_batch, + device_id=device_id) + return _normalize_prediction(out) + + return _predict + + def _build_predict_reduced_unrolled(self): + static_positions = self._reduced_static_positions + syndrome_positions = self._reduced_syndrome_positions + recipe_positions = self._reduced_recipe_positions + path_steps = self._reduced_path_steps + n = self._reduced_n_tensors def _predict(noise_probs: torch.Tensor, syndrome_tuple: tuple[torch.Tensor, ...]) -> torch.Tensor: - noise_stacked = torch.stack((1.0 - noise_probs, noise_probs), - dim=-1) ops: list[torch.Tensor] = [None] * n # type: ignore - for pos, arr in static_arrays.items(): + for pos, arr in static_positions.items(): ops[pos] = arr - for pos, arr in zip(syndrome_positions, syndrome_tuple): + for pos, syndrome_idx in syndrome_positions: + ops[pos] = syndrome_tuple[syndrome_idx] + for pos, arr in zip(recipe_positions, + self._precontract_reduced_noise(noise_probs)): ops[pos] = arr - for k, pos in enumerate(noise_pos_ordered): - ops[pos] = noise_stacked[k] for eq_str, idxs, sorted_idxs in path_steps: picked = [ops[i] for i in idxs] for i in sorted_idxs: ops.pop(i) ops.append(torch.einsum(eq_str, *picked)) out = ops[0] - return out / out.sum(dim=1, keepdim=True) + return _normalize_prediction(out) return _predict - def _build_predict_codegen(self): - """Codegen predict: partial-eval'd flat Python with named locals.""" - static_arrays = self._static_arrays - syndrome_positions = tuple(p for p, _t in self._syndrome_positions) - noise_pos_ordered = self._noise_pos_ordered - n = len(self._tensors_ref) - syndrome_tensors = list(self._syndrome_arrays) - codegen_fn = self._build_codegen_predict( - n, + def _build_predict_reduced_codegen(self): + static_arrays = dict(self._reduced_static_positions) + dynamic_positions = [ + (pos, f"_R{idx}") + for idx, pos in enumerate(self._reduced_recipe_positions) + ] + if self._dynamic_syndromes: + dynamic_positions.extend( + (pos, f"_S{sidx}") + for pos, sidx in self._reduced_syndrome_positions) + else: + for pos, sidx in self._reduced_syndrome_positions: + static_arrays[pos] = self._syndrome_arrays[sidx] + + codegen_fn = self._build_codegen_reduced_predict( + self._reduced_n_tensors, static_arrays, - syndrome_positions, - noise_pos_ordered, - self._path_steps, - syndrome_tensors, + tuple(dynamic_positions), + self._reduced_path_steps, + n_recipes=self._reduced_n_recipes, + n_syndromes=len(self._reduced_syndrome_positions), dynamic_syndromes=self._dynamic_syndromes, ) self._codegen_fn = codegen_fn @@ -561,17 +867,21 @@ def _build_predict_codegen(self): self._codegen_n_runtime = getattr(codegen_fn, "_n_runtime", 0) if self._dynamic_syndromes: - return codegen_fn - # Static mode bakes syndromes into the closure and returns a - # 1-arg callable; wrap to match the public 2-arg signature. - def _predict_static( - noise_probs: torch.Tensor, - syndrome_tuple: tuple[torch.Tensor, ...] = () - ) -> torch.Tensor: - return codegen_fn(noise_probs) + def _predict( + noise_probs: torch.Tensor, + syndrome_tuple: tuple[torch.Tensor, ...]) -> torch.Tensor: + return codegen_fn(self._precontract_reduced_noise(noise_probs), + syndrome_tuple) + else: + + def _predict( + noise_probs: torch.Tensor, + syndrome_tuple: tuple[torch.Tensor, ...] = () + ) -> torch.Tensor: + return codegen_fn(self._precontract_reduced_noise(noise_probs)) - return _predict_static + return _predict def _maybe_torch_compile(self, fn, *, kind: str): """Wrap ``fn`` with :func:`torch.compile` if requested. @@ -579,7 +889,7 @@ def _maybe_torch_compile(self, fn, *, kind: str): On any compile failure, warn and fall back to eager. ``kind`` is included in the warning to disambiguate predict vs loss. """ - if not self._use_torch_compile: + if not self._use_torch_compile or self._execute_mode == "cutensornet": return fn try: kwargs = self._torch_compile_kwargs() @@ -599,10 +909,7 @@ def _compile_loss(self) -> None: Two variants are produced: one accepting logits (sigmoid applied inside) and one accepting probabilities directly. """ - if self._execute_mode == "codegen": - logits_fn, probs_fn = self._build_loss_codegen() - else: - logits_fn, probs_fn = self._build_loss_wrapped() + logits_fn, probs_fn = self._build_loss_wrapped() self._loss_from_logits_fn = logits_fn self._loss_from_probs_fn = probs_fn @@ -611,57 +918,6 @@ def _compile_loss(self) -> None: self._compiled_loss_from_probs = self._maybe_torch_compile(probs_fn, kind="loss") - def _build_loss_codegen(self): - """Codegen loss: fuse the CE reduction into the contraction graph.""" - static_arrays = self._static_arrays - syndrome_positions = tuple(p for p, _t in self._syndrome_positions) - noise_pos_ordered = self._noise_pos_ordered - n = len(self._tensors_ref) - syndrome_tensors = list(self._syndrome_arrays) - - codegen_logits = self._build_codegen_loss( - n, - static_arrays, - syndrome_positions, - noise_pos_ordered, - self._path_steps, - syndrome_tensors, - obs_idx_true=self.obs_idx_true, - obs_idx_false=self.obs_idx_false, - dynamic_syndromes=self._dynamic_syndromes, - from_logits=True, - ) - codegen_probs = self._build_codegen_loss( - n, - static_arrays, - syndrome_positions, - noise_pos_ordered, - self._path_steps, - syndrome_tensors, - obs_idx_true=self.obs_idx_true, - obs_idx_false=self.obs_idx_false, - dynamic_syndromes=self._dynamic_syndromes, - from_logits=False, - ) - - if self._dynamic_syndromes: - return codegen_logits, codegen_probs - - # Static codegen bakes syndromes into the closure and returns a - # 1-arg callable; wrap to match the public 2-arg signature. - def _loss_from_logits_static( - logits: torch.Tensor, syndrome_tuple: tuple[torch.Tensor, ...] = () - ) -> torch.Tensor: - return codegen_logits(logits) - - def _loss_from_probs_static( - noise_probs: torch.Tensor, - syndrome_tuple: tuple[torch.Tensor, ...] = () - ) -> torch.Tensor: - return codegen_probs(noise_probs) - - return _loss_from_logits_static, _loss_from_probs_static - def _build_loss_wrapped(self): """opt_einsum / unrolled loss: wrap CE around ``self._predict_fn``.""" obs_t = self.obs_idx_true @@ -708,59 +964,28 @@ def _torch_compile_kwargs(self) -> dict[str, Any]: return kwargs @staticmethod - def _codegen_partial_eval(n, static_arrays, syndrome_positions, - noise_pos_ordered, path_steps, syndrome_tensors, - dynamic_syndromes: bool): - """Partial-evaluate ``path_steps``; return the codegen building blocks. - - Steps whose inputs are all static are evaluated eagerly under - ``torch.no_grad`` and become closure constants; the remaining - steps become source lines. - - Returns ``(runtime_lines, closure_vars, used_static, final_state, - n_folded)``: emitted source lines, name -> tensor map for the - function namespace, the subset of names actually referenced, the - single surviving state slot ``(name, is_dynamic, value_or_None)``, - and the count of folded steps. - """ + def _codegen_partial_eval_dynamic(n, static_arrays, dynamic_positions, + path_steps): static_positions = sorted(static_arrays.keys()) - noise_pos_set = set(noise_pos_ordered) - syn_pos_set = set(syndrome_positions) - # O(1) reverse lookup tables (avoid repeated list.index() inside - # the per-step loop below — was O(N^2) for large path lengths). - noise_pos_to_k = {pos: k for k, pos in enumerate(noise_pos_ordered)} - syn_pos_to_sidx = { - pos: sidx for sidx, pos in enumerate(syndrome_positions) - } + dynamic_names = {pos: name for pos, name in dynamic_positions} static_pos_to_sidx = { pos: sidx for sidx, pos in enumerate(static_positions) } - # state[pos] = (var_name, is_dynamic, concrete_value_or_None) state: list[tuple[str, bool, torch.Tensor | None]] = [] for pos in range(n): - if pos in noise_pos_set: - k = noise_pos_to_k[pos] - state.append((f"_n{k}", True, None)) - elif pos in syn_pos_set: - sidx = syn_pos_to_sidx[pos] - if dynamic_syndromes: - state.append((f"_S{sidx}", True, None)) - else: - state.append((f"_S{sidx}", False, syndrome_tensors[sidx])) + if pos in dynamic_names: + state.append((dynamic_names[pos], True, None)) else: sidx = static_pos_to_sidx[pos] state.append( (f"_C{sidx}", False, static_arrays[static_positions[sidx]])) - closure_vars: dict[str, torch.Tensor] = {} + closure_vars = { + f"_C{sidx}": static_arrays[pos] + for sidx, pos in enumerate(static_positions) + } runtime_lines: list[str] = [] - # Names that the emitted source actually references and that must - # be available in the closure namespace. We track this - # structurally as we go instead of re-parsing the source string, - # which is both faster and immune to lexical false positives / - # negatives (e.g. if an einsum equation contained an underscore). - used_static: set[str] = set() n_folded = 0 for step_idx, step in enumerate(path_steps): @@ -781,222 +1006,68 @@ def _codegen_partial_eval(n, static_arrays, syndrome_positions, n_folded += 1 else: arg_names = [p[0] for p in picked] - # Track which closure names this line will read. ``_n*`` - # names are header-built locals (from noise_probs) so - # they are *not* closure values; everything else is. - for name in arg_names: - if name.startswith(("_C", "_P")): - used_static.add(name) - elif name.startswith("_S") and not dynamic_syndromes: - used_static.add(name) runtime_lines.append( f" {out_name} = torch.einsum({eq_str!r}, " f"{', '.join(arg_names)})") state.append((out_name, True, None)) assert len(state) == 1 - for name in used_static: - if name in closure_vars: - continue - if name.startswith("_C"): - sidx = int(name[2:]) - closure_vars[name] = static_arrays[static_positions[sidx]] - elif name.startswith("_S"): # static-syndromes mode only - sidx = int(name[2:]) - closure_vars[name] = syndrome_tensors[sidx] - - return runtime_lines, closure_vars, used_static, state[0], n_folded - - @staticmethod - def _emit_noise_header(noise_pos_ordered, - transform: str = "identity") -> list[str]: - """Emit source lines materialising ``_n0 .. _n{K-1}``. - - ``transform="identity"`` treats the input as probabilities; - ``"sigmoid"`` treats it as logits and applies ``torch.sigmoid`` - first. A single ``(K, 2)`` stack is built and then sliced, which - keeps the autograd graph compact. - """ - lines: list[str] = [] - if transform == "sigmoid": - lines.append(" _p = torch.sigmoid(noise_probs)") - else: - lines.append(" _p = noise_probs") - lines.append(" _q = 1.0 - _p") - # One stack of shape (K, 2) instead of K stacks of shape (2,). - # ``dim=1`` makes ``_NS[k]`` a contiguous view of length 2. - lines.append(" _NS = torch.stack((_q, _p), dim=1)") - for k in range(len(noise_pos_ordered)): - lines.append(f" _n{k} = _NS[{k}]") - return lines - - @staticmethod - def _emit_syndrome_header(syndrome_positions, - dynamic_syndromes: bool) -> list[str]: - """Emit source lines binding ``_S0 .. _S{S-1}`` to runtime - ``syndromes`` arguments; empty in static mode.""" - if not dynamic_syndromes: - return [] - return [ - f" _S{sidx} = syndromes[{sidx}]" - for sidx in range(len(syndrome_positions)) - ] + return runtime_lines, closure_vars, state[0], n_folded @classmethod - def _build_codegen_predict(cls, - n, - static_arrays, - syndrome_positions, - noise_pos_ordered, - path_steps, - syndrome_tensors, - dynamic_syndromes: bool = True): - """Generate ``_predict(noise_probs[, syndromes]) -> (shots, 2)``. - - With ``dynamic_syndromes=False`` syndromes are folded into the - closure, which maximises partial evaluation but forces a rebuild - on every :meth:`update_dataset` call. With ``True`` (default) - syndromes stay runtime arguments and dataset swaps are free. - """ - runtime_lines, closure_vars, _used, final_state, n_folded = ( - cls._codegen_partial_eval( + def _build_codegen_reduced_predict(cls, + n, + static_arrays, + dynamic_positions, + path_steps, + n_recipes: int, + n_syndromes: int, + dynamic_syndromes: bool = True): + runtime_lines, closure_vars, final_state, n_folded = ( + cls._codegen_partial_eval_dynamic( n, static_arrays, - syndrome_positions, - noise_pos_ordered, + dynamic_positions, path_steps, - syndrome_tensors, - dynamic_syndromes, )) final_name, is_final_dyn, final_value = final_state fully_static = not is_final_dyn body: list[str] = [] if dynamic_syndromes: - body.append("def _predict(noise_probs, syndromes):") + body.append("def _predict(recipe_arrays, syndromes):") else: - body.append("def _predict(noise_probs):") + body.append("def _predict(recipe_arrays):") if fully_static: - # Degenerate case: contraction didn't depend on noise (or any - # other dynamic input). Pre-normalise the constant and - # return it unchanged on every call. with torch.no_grad(): - normed = final_value / final_value.sum(dim=1, keepdim=True) + normed = _normalize_prediction(final_value) closure_vars["_FINAL"] = normed body.append(" return _FINAL") runtime_lines = [] else: - body.extend( - cls._emit_noise_header(noise_pos_ordered, transform="identity")) - body.extend( - cls._emit_syndrome_header(syndrome_positions, - dynamic_syndromes)) + for k in range(n_recipes): + body.append(f" _R{k} = recipe_arrays[{k}]") + if dynamic_syndromes: + for sidx in range(n_syndromes): + body.append(f" _S{sidx} = syndromes[{sidx}]") body.extend(runtime_lines) body.append(f" _out = {final_name}") - body.append(" return _out / _out.sum(dim=1, keepdim=True)") + body.append(" return _normalize_prediction(_out)") return cls._compile_codegen_source(body, closure_vars, n_folded, len(runtime_lines), "predict") - @classmethod - def _build_codegen_loss(cls, - n, - static_arrays, - syndrome_positions, - noise_pos_ordered, - path_steps, - syndrome_tensors, - obs_idx_true: torch.Tensor, - obs_idx_false: torch.Tensor, - dynamic_syndromes: bool = True, - from_logits: bool = True): - """Generate a fused ``(input, syndromes) -> scalar`` loss callable. - - Pipes the contraction output straight into the cross-entropy - reduction so the whole pipeline (optional sigmoid, contraction, - normalisation, cross-entropy) is a single autograd graph. - - Args: - from_logits: If ``True`` (default), apply ``torch.sigmoid`` to - the input; if ``False``, the input must already be in - ``[0, 1]``. - """ - runtime_lines, closure_vars, _used, final_state, n_folded = ( - cls._codegen_partial_eval( - n, - static_arrays, - syndrome_positions, - noise_pos_ordered, - path_steps, - syndrome_tensors, - dynamic_syndromes, - )) - final_name, is_final_dyn, final_value = final_state - fully_static = not is_final_dyn - - closure_vars["_OBS_T"] = obs_idx_true - closure_vars["_OBS_F"] = obs_idx_false - - body: list[str] = [] - if dynamic_syndromes: - body.append("def _loss(noise_probs, syndromes):") - else: - body.append("def _loss(noise_probs):") - - if fully_static: - # The contraction is a constant; the loss it produces is - # also a constant. We still need the gradient wrt - # noise_probs to be zero (a real-valued zero tensor with a - # graph edge), so emit a no-op multiplication by - # noise_probs.sum() * 0 to keep autograd happy. - with torch.no_grad(): - normed = final_value / final_value.sum(dim=1, keepdim=True) - # Compute the loss eagerly; we can't fold it because - # autograd needs a path back to noise_probs. - ce = ( - -torch.log(_clamp_log_input(normed[obs_idx_true, 1])).sum() - - - torch.log(_clamp_log_input(normed[obs_idx_false, 0])).sum()) - closure_vars["_LOSS"] = ce - body.append(" return _LOSS + 0.0 * noise_probs.sum()") - runtime_lines = [] - else: - transform = "sigmoid" if from_logits else "identity" - body.extend(cls._emit_noise_header(noise_pos_ordered, transform)) - body.extend( - cls._emit_syndrome_header(syndrome_positions, - dynamic_syndromes)) - body.extend(runtime_lines) - # Fused cross-entropy that skips the explicit - # ``_preds = _out / _out.sum(dim=1, keepdim=True)`` step. - # OBS_T and OBS_F partition the batch (every row is in exactly - # one), so: - # -log(p_T[:,1]).sum() - log(p_F[:,0]).sum() - # = log(Z).sum() - log(_out[OBS_T,1]).sum() - # - log(_out[OBS_F,0]).sum() - # where Z = _out[:,0] + _out[:,1]. Saves one (shots, 2) - # division + materialisation and the corresponding backward - # nodes -- ~2-5% per-step on CPU at d=3/r=3. - body.append(f" _out = {final_name}") - body.append(" _z0 = _out[:, 0]") - body.append(" _z1 = _out[:, 1]") - body.append(" _eps = torch.finfo(_z0.dtype).tiny") - body.append( - " return (torch.log((_z0 + _z1).clamp_min(_eps)).sum() " - "- torch.log(_z1[_OBS_T].clamp_min(_eps)).sum() " - "- torch.log(_z0[_OBS_F].clamp_min(_eps)).sum())") - - return cls._compile_codegen_source(body, closure_vars, n_folded, - len(runtime_lines), "loss") - @staticmethod def _compile_codegen_source(body: list[str], closure_vars: dict[str, torch.Tensor], n_folded: int, n_runtime: int, kind: str): """Compile the assembled function source and return the callable.""" source = "\n".join(body) - ns: dict[str, Any] = {"torch": torch} + ns: dict[str, Any] = { + "torch": torch, + "_normalize_prediction": _normalize_prediction, + } ns.update(closure_vars) fn_name = "_loss" if kind == "loss" else "_predict" exec(compile(source, f"", "exec"), ns) @@ -1067,7 +1138,7 @@ def _update_data(self, should use :meth:`update_dataset` instead. """ # Patch syndrome tensor data in the quimb TN in place; the - # cotengra path is invalidated below if any shape changed. + # contraction path is invalidated below if any shape changed. for i, tag in enumerate(self._syndrome_tags): t = self.syndrome_tn.tensors[next( iter(self.syndrome_tn.tag_map[tag]))] @@ -1100,7 +1171,7 @@ def _update_data(self, # Shape change: cached path / codegen / oe expression / compile # guards are all stale. Drop the path and rebuild from scratch. # Shapes unchanged: dynamic codegen and the unrolled / - # opt_einsum paths read syndromes per call — refreshing the + # opt_einsum paths read syndromes per call - refreshing the # cached tuple is enough. Static codegen baked the old tensors # into the closure and still needs a full rebuild. shape_changed = new_shapes_tuple != self._syndrome_shapes @@ -1154,7 +1225,7 @@ def optimize_path(self, optimize: Any = None, batch_size: int = -1) -> Any: Always routes through :meth:`TensorNetwork.contraction_info` so the resulting path is compatible with :mod:`opt_einsum` and manual unrolling -- unlike :meth:`TensorNetworkDecoder.optimize_path`, - which defaults to a cuTensorNet-only path. + which may return a path not usable by these torch-backed modes. ``batch_size`` is part of the parent ``TensorNetworkDecoder`` signature (which rebuilds its TN around a fake batch); on the @@ -1163,14 +1234,10 @@ def optimize_path(self, optimize: Any = None, batch_size: int = -1) -> Any: ignored. Kept for Liskov substitution with the parent. """ del batch_size - info = self.full_tn.contraction_info( - output_inds=("batch_index", self.logical_obs_inds[0]), - optimize=optimize if optimize is not None else "auto", - ) - self.path_batch = info.path + self.path_batch = optimize if optimize is not None else "auto" self.slicing_batch = tuple() self._snapshot_arrays_and_eq() - return info + return self._reduced_info def make_compiled_step(optimizer: NMOptimizer, logits: torch.Tensor, @@ -1203,3 +1270,184 @@ def _step(): return loss return _step + + +def make_batch_sliced_step( + optimizer: NMOptimizer, + logits: torch.Tensor, + torch_optimizer: torch.optim.Optimizer, + progress_callback: Callable[[int, int, float], None] | None = None): + """Build a step callable that streams a larger batch through slices. + + The optimizer should be constructed with the desired slice-sized + dataset. The returned callable accepts a larger raw syndrome batch, + repeatedly updates the optimizer with fixed-shape slices, accumulates + gradients, and performs one ``torch_optimizer.step()``. This keeps + path finding and contraction execution tied to the smaller batch + shape while preserving the summed cross-entropy objective over all + supplied shots. The optimizer's live dataset is left at the final + slice after the step. + + Args: + optimizer: An :class:`NMOptimizer` constructed with the desired + batch-slice size. + logits: Trainable 1-D tensor of length ``len(optimizer.error_inds)`` + with ``requires_grad=True``. + torch_optimizer: A ``torch.optim`` instance owning ``logits``. + progress_callback: Optional callback invoked as + ``callback(done_chunks, total_chunks, elapsed_seconds)``. + + Returns: + A callable ``step(syndrome_data, observable_flips)`` that runs one + optimizer step and returns the detached summed loss. + """ + if not optimizer._dynamic_syndromes: + raise ValueError("batch-sliced steps require dynamic_syndromes=True.") + + batch_slice_size = int(optimizer._batch_size) + if batch_slice_size <= 0: + raise ValueError("optimizer batch size must be positive.") + + def _step(syndrome_data: npt.NDArray[Any], + observable_flips: npt.NDArray[Any]) -> torch.Tensor: + syndrome_arr = np.asarray(syndrome_data) + flips_arr = np.asarray(observable_flips, dtype=bool).reshape(-1) + if syndrome_arr.ndim != 2: + raise ValueError("syndrome_data must have shape (shots, checks).") + if syndrome_arr.shape[0] != flips_arr.shape[0]: + raise ValueError( + "syndrome_data and observable_flips length mismatch.") + if syndrome_arr.shape[0] == 0: + raise ValueError("syndrome_data must contain at least one shot.") + + torch_optimizer.zero_grad(set_to_none=True) + total_loss = torch.zeros((), dtype=logits.dtype, device=logits.device) + num_chunks = ((syndrome_arr.shape[0] + batch_slice_size - 1) // + batch_slice_size) + wall_start = None + if progress_callback is not None: + import time + wall_start = time.perf_counter() + + for chunk_idx, start in enumerate(range(0, syndrome_arr.shape[0], + batch_slice_size), + start=1): + stop = min(start + batch_slice_size, syndrome_arr.shape[0]) + valid = stop - start + chunk_syn = syndrome_arr[start:stop] + chunk_flips = flips_arr[start:stop] + if valid < batch_slice_size: + padded_syn = np.zeros((batch_slice_size, syndrome_arr.shape[1]), + dtype=syndrome_arr.dtype) + padded_flips = np.zeros((batch_slice_size,), dtype=bool) + padded_syn[:valid] = chunk_syn + padded_flips[:valid] = chunk_flips + chunk_syn = padded_syn + update_flips = padded_flips + else: + update_flips = chunk_flips + + optimizer.update_dataset(chunk_syn, + update_flips, + enforce_shape=True) + preds = optimizer._compiled_predict( + torch.sigmoid(logits), optimizer.current_syndrome_args()) + preds = preds[:valid] + flips_t = torch.as_tensor(chunk_flips, + dtype=torch.bool, + device=preds.device) + loss = (-torch.log(_clamp_log_input(preds[flips_t, 1])).sum() - + torch.log(_clamp_log_input(preds[~flips_t, 0])).sum()) + loss.backward() + total_loss = total_loss + loss.detach() + + if progress_callback is not None: + elapsed = (0.0 if wall_start is None else time.perf_counter() - + wall_start) + progress_callback(chunk_idx, num_chunks, elapsed) + + torch_optimizer.step() + return total_loss + + return _step + + +def make_training_step(H: npt.NDArray[Any], + logical_obs: npt.NDArray[Any], + noise_model: list[float], + syndrome_data: npt.NDArray[Any], + observable_flips: npt.NDArray[Any], + logits: torch.Tensor, + torch_optimizer: torch.optim.Optimizer, + batch_slicing: bool | Literal["auto"] = "auto", + batch_slice_size: int | None = None, + progress_callback: Callable[[int, int, float], None] | + None = None, + **optimizer_kwargs: Any) -> tuple[NMOptimizer, Any]: + """Construct an :class:`NMOptimizer` and matching training step. + + Args: + H: Parity check matrix. + logical_obs: Logical observable matrix. + noise_model: Initial per-error probabilities. + syndrome_data: Full syndrome batch, shape ``(shots, checks)``. + observable_flips: Observable flips for the full batch. + logits: Trainable logit tensor owned by ``torch_optimizer``. + torch_optimizer: Optimizer for ``logits``. + batch_slicing: ``False`` uses the full batch. ``"auto"`` or + ``True`` constructs a slice-sized optimizer and streams the + full batch through it with gradient accumulation. + batch_slice_size: Optional explicit slice size. Defaults to ``1`` + when batch slicing is enabled. + progress_callback: Optional callback forwarded to the batch-sliced + step as ``callback(done_chunks, total_chunks, elapsed_seconds)``. + **optimizer_kwargs: Forwarded to :class:`NMOptimizer`. + + Returns: + ``(optimizer, step_fn)``. For full-batch training ``step_fn`` is + no-arg. For batch-sliced training it accepts optional + ``(syndrome_data, observable_flips)`` and defaults to the original + full batch supplied here. + """ + if batch_slicing not in (False, True, "auto"): + raise ValueError("batch_slicing must be False, True, or 'auto'.") + + use_batch_slicing = batch_slicing in (True, "auto") + syndrome_arr = np.asarray(syndrome_data) + flips_arr = np.asarray(observable_flips, dtype=bool).reshape(-1) + if syndrome_arr.ndim != 2: + raise ValueError("syndrome_data must have shape (shots, checks).") + if syndrome_arr.shape[0] != flips_arr.shape[0]: + raise ValueError("syndrome_data and observable_flips length mismatch.") + if syndrome_arr.shape[0] == 0: + raise ValueError("syndrome_data must contain at least one shot.") + + if not use_batch_slicing: + optimizer = NMOptimizer(H, logical_obs, noise_model, syndrome_arr, + flips_arr, **optimizer_kwargs) + return optimizer, make_compiled_step(optimizer, logits, torch_optimizer) + + slice_size = 1 if batch_slice_size is None else int(batch_slice_size) + if slice_size <= 0: + raise ValueError("batch_slice_size must be a positive integer.") + slice_size = min(slice_size, syndrome_arr.shape[0]) + optimizer = NMOptimizer(H, logical_obs, noise_model, + syndrome_arr[:slice_size], flips_arr[:slice_size], + **optimizer_kwargs) + sliced_step = make_batch_sliced_step(optimizer, logits, torch_optimizer, + progress_callback) + + def _step( + new_syndrome_data: npt.NDArray[Any] | None = None, + new_observable_flips: npt.NDArray[Any] | None = None + ) -> torch.Tensor: + if new_syndrome_data is None: + if new_observable_flips is not None: + raise ValueError( + "observable_flips provided without syndrome_data.") + return sliced_step(syndrome_arr, flips_arr) + if new_observable_flips is None: + raise ValueError("observable_flips is required with syndrome_data.") + return sliced_step(new_syndrome_data, new_observable_flips) + + return optimizer, _step diff --git a/libs/qec/python/tests/test_nm_optimizer.py b/libs/qec/python/tests/test_nm_optimizer.py index 97dceaba1..034f6255a 100644 --- a/libs/qec/python/tests/test_nm_optimizer.py +++ b/libs/qec/python/tests/test_nm_optimizer.py @@ -22,9 +22,13 @@ "torch", reason="torch not installed; skipping TN noise-learning tests") if sys.version_info >= (3, 11): + from cudaq_qec.plugins.decoders.tensor_network_utils import ( + nm_optimizer as nm_optimizer_mod,) from cudaq_qec.plugins.decoders.tensor_network_utils.nm_optimizer import ( NMOptimizer, + make_batch_sliced_step, make_compiled_step, + make_training_step, remap_eq_to_ascii, ) @@ -111,6 +115,32 @@ def _naive_cross_entropy(opt: "NMOptimizer") -> torch.Tensor: torch.log(preds[obs_f, 0]).sum()) +def _full_network_prediction(opt: "NMOptimizer") -> torch.Tensor: + arrays = [] + noise_ids = {id(t) for t in opt.noise_model.tensors} + noise_stacked = torch.stack( + (1.0 - opt.noise_params[0], opt.noise_params[0]), dim=-1) + noise_pos_by_error = { + id(t): k for k, t in enumerate(opt.noise_model.tensors) + } + for tensor in opt.full_tn.tensors: + if id(tensor) in noise_ids: + arrays.append(noise_stacked[noise_pos_by_error[id(tensor)]]) + elif isinstance(tensor.data, torch.Tensor): + arrays.append(tensor.data.detach().to( + device=opt.torch_device, dtype=opt.noise_params[0].dtype)) + else: + arrays.append( + torch.as_tensor(np.asarray(tensor.data), + dtype=opt.noise_params[0].dtype, + device=opt.torch_device)) + out = torch.einsum( + opt.full_tn.get_equation(output_inds=("batch_index", + opt.logical_obs_inds[0])), + *arrays) + return out / out.sum(dim=1, keepdim=True) + + # -- construction ------------------------------------------------------------ @@ -123,6 +153,8 @@ def test_construction_basic(device): num_shots=8, rng=np.random.default_rng(0)) opt = _make_opt(H, logical, priors, syn, flips, device=device) + assert opt.contractor_config.contractor_name == "oe_torch_compiled" + assert opt.contractor_config.backend == "torch" assert opt._batch_size == 8 assert opt._noise_probs.requires_grad assert len(opt.noise_params) == 1 @@ -171,6 +203,23 @@ def test_invalid_dtype_rejected(device): dtype="float16") +def test_cutensornet_execute_requires_cuda(): + H, logical, priors = _simple_repetition_code() + syn, flips = _sample_synthetic_dataset(H, + logical, + priors, + num_shots=4, + rng=np.random.default_rng(11)) + with pytest.raises(ValueError, match="requires a CUDA device"): + _make_opt(H, + logical, + priors, + syn, + flips, + device="cpu", + execute="cutensornet") + + # -- forward pass / gradient ------------------------------------------------- @@ -329,6 +378,121 @@ def test_loss_fn_from_logits_and_probs(device, execute): assert torch.allclose(v_probs, v_self, atol=1e-8, rtol=1e-8) +@pytest.mark.parametrize("device", _device_params()) +@pytest.mark.parametrize("execute", _EXECUTE_MODES) +def test_reduced_path_matches_full_network_reference(device, execute): + rng = np.random.default_rng(26) + H, logical = _nondegenerate_code() + init_priors = [0.2, 0.3, 0.4] + syn, flips = _sample_synthetic_dataset(H, + logical, [0.1, 0.15, 0.25], + num_shots=24, + rng=rng) + opt = _make_opt(H, + logical, + init_priors, + syn, + flips, + device=device, + dtype="float64", + execute=execute) + with torch.no_grad(): + p_full = _full_network_prediction(opt) + p_reduced = opt.decoder_prediction() + loss_full = (-torch.log(p_full[opt.obs_idx_true, 1]).sum() - + torch.log(p_full[opt.obs_idx_false, 0]).sum()) + loss_reduced = opt.cross_entropy_loss() + assert torch.allclose(p_full, p_reduced, atol=1e-7, rtol=1e-7) + assert torch.allclose(loss_full, loss_reduced, atol=1e-7, rtol=1e-7) + + opt.noise_params[0].grad = None + loss = opt.cross_entropy_loss() + loss.backward() + grad = opt.noise_params[0].grad + assert grad is not None + assert torch.isfinite(grad).all() + assert torch.any(grad != 0.0) + + +@pytest.mark.parametrize("device", _device_params()) +def test_reduced_path_matches_full_network_reference_static_codegen(device): + rng = np.random.default_rng(27) + H, logical = _nondegenerate_code() + init_priors = [0.2, 0.3, 0.4] + syn, flips = _sample_synthetic_dataset(H, + logical, [0.1, 0.15, 0.25], + num_shots=24, + rng=rng) + opt = _make_opt(H, + logical, + init_priors, + syn, + flips, + device=device, + dtype="float64", + execute="codegen", + dynamic_syndromes=False) + with torch.no_grad(): + full = _full_network_prediction(opt) + assert torch.allclose(full, + opt.decoder_prediction(), + atol=1e-7, + rtol=1e-7) + loss_full = (-torch.log(full[opt.obs_idx_true, 1]).sum() - + torch.log(full[opt.obs_idx_false, 0]).sum()) + assert torch.allclose(loss_full, + opt.cross_entropy_loss(), + atol=1e-7, + rtol=1e-7) + + +@pytest.mark.skipif(not _gpu_available(), reason="CUDA not available") +def test_cutensornet_execute_matches_opt_einsum_and_backpropagates(): + """cuTN can execute the reduced torch graph with autograd.""" + pytest.importorskip("cuquantum") + rng = np.random.default_rng(31415) + H, logical = _nondegenerate_code() + init_priors = [0.2, 0.3, 0.4] + syn, flips = _sample_synthetic_dataset(H, + logical, [0.1, 0.15, 0.25], + num_shots=12, + rng=rng) + ref = _make_opt(H, + logical, + init_priors, + syn, + flips, + device="cuda", + dtype="float64", + execute="opt_einsum") + cutn = _make_opt(H, + logical, + init_priors, + syn, + flips, + device="cuda", + dtype="float64", + execute="cutensornet") + + with torch.no_grad(): + p_ref = ref.decoder_prediction() + p_cutn = cutn.decoder_prediction() + loss_ref = ref.cross_entropy_loss() + loss_cutn = cutn.cross_entropy_loss() + + assert torch.allclose(p_cutn, p_ref, atol=1e-10, rtol=1e-10) + assert torch.allclose(loss_cutn, loss_ref, atol=1e-10, rtol=1e-10) + assert cutn.contractor_config.contractor_name == "cutensornet" + + cutn.noise_params[0].grad = None + loss = cutn.cross_entropy_loss() + loss.backward() + grad = cutn.noise_params[0].grad + assert grad is not None + assert torch.isfinite(grad).all() + assert torch.any(grad != 0.0) + + # -- numerical guards -------------------------------------------------------- @@ -855,9 +1019,138 @@ def test_cpu_gpu_parity_forward(): np.testing.assert_allclose(p_cpu, p_gpu, atol=1e-4, rtol=1e-4) +def test_default_torch_path_prefers_rank_safe_candidate(monkeypatch): + """Path selection avoids steps PyTorch cannot materialize.""" + + class _Info: + + def __init__(self, largest, cost, eq): + self.largest_intermediate = largest + self.opt_cost = cost + self.contraction_list = [(None, None, eq)] + + def _fake_contract_path(eq, *shapes, optimize=None, **kwargs): + del eq, shapes, kwargs + if optimize == "greedy": + unsafe_eq = "abcdefghijklmnopqrstuvwxyz,ABCDEFGHIJKLMNOPQRSTUVWXYZ" \ + "->abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + return ["unsafe"], _Info(1.0, 1.0, unsafe_eq) + return ["safe"], _Info(10.0, 10.0, "ab,bc->ac") + + monkeypatch.setattr(nm_optimizer_mod.oe, "contract_path", + _fake_contract_path) + path, info = nm_optimizer_mod._select_default_torch_path( + "ab,bc->ac", ((2, 2), (2, 2))) + + assert path == ["safe"] + assert nm_optimizer_mod._path_max_einsum_rank(info) <= 25 + + +# -- batch-sliced training --------------------------------------------------- + + +@pytest.mark.parametrize("execute", _EXECUTE_MODES) +def test_batch_sliced_step_matches_full_step_with_tail(execute): + """Batch-sliced accumulation matches the full summed objective.""" + rng = np.random.default_rng(5150) + H, logical = _nondegenerate_code() + true_priors = [0.03, 0.12, 0.08] + syn, flips = _sample_synthetic_dataset(H, + logical, + true_priors, + num_shots=10, + rng=rng) + init_priors = [0.10] * H.shape[1] + full_opt = _make_opt(H, + logical, + init_priors, + syn, + flips, + device="cpu", + dtype="float64", + execute=execute) + slice_size = 4 + sliced_opt = _make_opt(H, + logical, + init_priors, + syn[:slice_size], + flips[:slice_size], + device="cpu", + dtype="float64", + execute=execute) + logits_full = torch.logit(full_opt.noise_params[0].detach()).clone() + logits_full.requires_grad_(True) + logits_sliced = torch.logit(sliced_opt.noise_params[0].detach()).clone() + logits_sliced.requires_grad_(True) + full_torch_opt = torch.optim.SGD([logits_full], lr=0.01) + sliced_torch_opt = torch.optim.SGD([logits_sliced], lr=0.01) + + full_loss = make_compiled_step(full_opt, logits_full, full_torch_opt)() + sliced_loss = make_batch_sliced_step(sliced_opt, logits_sliced, + sliced_torch_opt)(syn, flips) + + assert torch.allclose(full_loss.detach(), + sliced_loss, + atol=1e-10, + rtol=1e-10) + assert torch.allclose(logits_full.detach(), + logits_sliced.detach(), + atol=1e-10, + rtol=1e-10) + + # -- truth-data convergence -------------------------------------------------- +def test_make_training_step_auto_matches_full_batch_step(): + """Factory hides slice construction while preserving the objective.""" + rng = np.random.default_rng(8675309) + H, logical = _nondegenerate_code() + true_priors = [0.03, 0.12, 0.08] + syn, flips = _sample_synthetic_dataset(H, + logical, + true_priors, + num_shots=10, + rng=rng) + init_priors = [0.10] * H.shape[1] + + full_opt = _make_opt(H, + logical, + init_priors, + syn, + flips, + device="cpu", + dtype="float64", + execute="opt_einsum") + logits_full = torch.logit(full_opt.noise_params[0].detach()).clone() + logits_full.requires_grad_(True) + full_torch_opt = torch.optim.SGD([logits_full], lr=0.01) + full_loss = make_compiled_step(full_opt, logits_full, full_torch_opt)() + + logits_auto = torch.logit(full_opt.noise_params[0].detach()).clone() + logits_auto.requires_grad_(True) + auto_torch_opt = torch.optim.SGD([logits_auto], lr=0.01) + auto_opt, auto_step = make_training_step(H, + logical, + init_priors, + syn, + flips, + logits_auto, + auto_torch_opt, + batch_slicing="auto", + device="cpu", + dtype="float64", + execute="opt_einsum") + assert auto_opt._batch_size == 1 + auto_loss = auto_step() + + assert torch.allclose(full_loss.detach(), auto_loss, atol=1e-10, rtol=1e-10) + assert torch.allclose(logits_full.detach(), + logits_auto.detach(), + atol=1e-10, + rtol=1e-10) + + @pytest.mark.parametrize("device", _device_params()) def test_recovers_true_priors_within_tol(device): """Fitted priors converge to the Bernoulli rates that sampled the data.""" diff --git a/libs/qec/python/tests/test_tensor_network_decoder.py b/libs/qec/python/tests/test_tensor_network_decoder.py index 4fae06c4a..6ce5031dd 100644 --- a/libs/qec/python/tests/test_tensor_network_decoder.py +++ b/libs/qec/python/tests/test_tensor_network_decoder.py @@ -23,7 +23,8 @@ prepare_syndrome_data_batch, tensor_network_from_syndrome_batch, tensor_network_from_logical_observable) from cudaq_qec.plugins.decoders.tensor_network_utils.contractors import ( - optimize_path, cutn_contractor, ContractorConfig, contractor) + optimize_path, cutn_contractor, ContractorConfig, contractor, + oe_torch_contractor, oe_torch_compiled_contractor) from cudaq_qec.plugins.decoders.tensor_network_utils.noise_models import factorized_noise_model, error_pairs_noise_model pytestmark = pytest.mark.skipif(sys.version_info < (3, 11), @@ -528,40 +529,28 @@ def test_error_pairs_noise_model_default_tags(): assert "NOISE" in t.tags -def test_valid_numpy_cpu(): - cfg = ContractorConfig("numpy", "numpy", "cpu") - assert cfg.contractor_name == "numpy" - assert cfg.backend == "numpy" - assert cfg.device == "cpu" +@pytest.mark.parametrize( + "contractor_name,backend,device,expected_contractor", + [ + ("numpy", "numpy", "cpu", contractor), + ("torch", "torch", "cpu", contractor), + ("torch", "torch", "cuda:0", contractor), + ("oe_torch", "torch", "cpu", oe_torch_contractor), + ("oe_torch", "torch", "cuda:0", oe_torch_contractor), + ("oe_torch_compiled", "torch", "cpu", oe_torch_compiled_contractor), + ("oe_torch_compiled", "torch", "cuda:0", oe_torch_compiled_contractor), + ("cutensornet", "numpy", "cuda", cutn_contractor), + ("cutensornet", "torch", "cuda", cutn_contractor), + ], +) +def test_valid_contractor_configs(contractor_name, backend, device, + expected_contractor): + cfg = ContractorConfig(contractor_name, backend, device) + assert cfg.contractor_name == contractor_name + assert cfg.backend == backend + assert cfg.device == device assert cfg.device_id == 0 - assert cfg.contractor is contractor - - -def test_valid_torch_cpu(): - cfg = ContractorConfig("torch", "torch", "cpu") - assert cfg.contractor_name == "torch" - assert cfg.backend == "torch" - assert cfg.device == "cpu" - assert cfg.device_id == 0 - assert cfg.contractor is contractor - - -def test_valid_cutensornet_numpy_cuda(): - cfg = ContractorConfig("cutensornet", "numpy", "cuda") - assert cfg.contractor_name == "cutensornet" - assert cfg.backend == "numpy" - assert cfg.device == "cuda" - assert cfg.device_id == 0 - assert cfg.contractor is cutn_contractor - - -def test_valid_cutensornet_torch_cuda(): - cfg = ContractorConfig("cutensornet", "torch", "cuda") - assert cfg.contractor_name == "cutensornet" - assert cfg.backend == "torch" - assert cfg.device == "cuda" - assert cfg.device_id == 0 - assert cfg.contractor is cutn_contractor + assert cfg.contractor is expected_contractor def test_cuda_device_id_parsing(): diff --git a/scripts/benchmark_nmoptimizer_d5r5_nico_dem.py b/scripts/benchmark_nmoptimizer_d5r5_nico_dem.py new file mode 100644 index 000000000..e6d90150b --- /dev/null +++ b/scripts/benchmark_nmoptimizer_d5r5_nico_dem.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python3 +# ============================================================================ # +# Copyright (c) 2026 NVIDIA Corporation & Affiliates. +# All rights reserved. +# +# This source code and the accompanying materials are made available under +# the terms of the Apache License 2.0 which accompanies this distribution. +# ============================================================================ # +"""NMOptimizer d=5/r=5 memory benchmark for large surface-code DEMs. + +This script is meant for the d=5/r=5 scalability discussion where the DEM has +roughly the same shape as Nico's case: 120 detectors and 1677 error mechanisms. +It can either generate a Stim rotated-memory surface-code DEM or load an exact +``.dem`` file. + +The script has two modes: + * path diagnostics: construct NMOptimizer and print reduced-path largest + intermediate / opt cost for one or more batch sizes. + * optional execution: run a tiny Adam step to expose actual forward/backward + memory failures. Use ``--run-step`` for this because the hard DEM can OOM. + +Example: + python3 scripts/benchmark_nmoptimizer_d5r5_nico_dem.py \ + --device cuda:0 --dtype float64 --execute all \ + --batch-sizes 1 2 4 1000 + +With an exact DEM file: + python3 scripts/benchmark_nmoptimizer_d5r5_nico_dem.py \ + --dem-path surface_code_bZ_d5_r05_center_5_5.dem \ + --batch-sizes 1 2 4 1000 +""" + +from __future__ import annotations + +import argparse +import contextlib +import gc +import math +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +try: + import stim +except ImportError as exc: # pragma: no cover + raise SystemExit("stim is required for this benchmark") from exc + +try: + from beliefmatching.belief_matching import detector_error_model_to_check_matrices +except ImportError as exc: # pragma: no cover + raise SystemExit("beliefmatching is required to convert DEMs") from exc + +from cudaq_qec import NMOptimizer, make_training_step + +EXECUTE_MODES = ("codegen", "unrolled", "opt_einsum", "cutensornet") + + +@dataclass +class Problem: + H: np.ndarray + logical: np.ndarray + priors: np.ndarray + circuit: stim.Circuit | None + dem: stim.DetectorErrorModel + + +@contextlib.contextmanager +def timed(label: str): + start = time.perf_counter() + print( + f"[START] {label}: {datetime.now(timezone.utc).isoformat(timespec='seconds')}", + flush=True) + try: + yield + except Exception: + elapsed = time.perf_counter() - start + print( + f"[FAILED] {label}: {datetime.now(timezone.utc).isoformat(timespec='seconds')} " + f"(elapsed {elapsed:.2f}s)", + flush=True) + raise + else: + elapsed = time.perf_counter() - start + print( + f"[END] {label}: {datetime.now(timezone.utc).isoformat(timespec='seconds')} " + f"(elapsed {elapsed:.2f}s)", + flush=True) + + +def parse_detector_error_model(dem: stim.DetectorErrorModel): + matrices = detector_error_model_to_check_matrices(dem) + H = np.zeros(matrices.check_matrix.shape, dtype=np.float64) + matrices.check_matrix.astype(np.float64).toarray(out=H) + logical = np.zeros(matrices.observables_matrix.shape, dtype=np.float64) + matrices.observables_matrix.astype(np.float64).toarray(out=logical) + priors = np.array([float(p) for p in matrices.priors], dtype=np.float64) + return H, logical, priors + + +def make_surface_code_circuit(args: argparse.Namespace) -> stim.Circuit: + noise_kwargs = { + "after_clifford_depolarization": args.after_clifford_p, + "after_reset_flip_probability": args.after_reset_p, + "before_measure_flip_probability": args.before_measure_p, + "before_round_data_depolarization": args.before_round_data_p, + } + noise_kwargs = { + k: v for k, v in noise_kwargs.items() if v is not None and v > 0.0 + } + return stim.Circuit.generated( + args.code, + distance=args.distance, + rounds=args.rounds, + **noise_kwargs, + ) + + +def build_problem(args: argparse.Namespace) -> Problem: + if args.dem_path is not None: + dem_text = Path(args.dem_path).read_text(encoding="utf-8") + dem = stim.DetectorErrorModel(dem_text) + circuit = None + else: + circuit = make_surface_code_circuit(args) + dem = circuit.detector_error_model(decompose_errors=True) + H, logical, priors = parse_detector_error_model(dem) + return Problem(H=H, + logical=logical, + priors=priors, + circuit=circuit, + dem=dem) + + +def make_circuit_sampler(circuit: stim.Circuit, seed: int): + try: + return circuit.compile_detector_sampler(seed=seed) + except TypeError: + return circuit.compile_detector_sampler() + + +def sample_from_circuit(circuit: stim.Circuit, shots: int, seed: int): + sampler = make_circuit_sampler(circuit, seed) + dets, obs = sampler.sample(shots, separate_observables=True) + return dets.astype(np.float64), obs.reshape(-1).astype(bool) + + +def sample_from_dem(dem: stim.DetectorErrorModel, shots: int, seed: int): + try: + sampler = dem.compile_sampler(seed=seed) + except TypeError: + sampler = dem.compile_sampler() + + # Stim versions have varied here. Try the more descriptive API first, + # then fall back to tuple shape inspection. + try: + dets, obs = sampler.sample(shots, separate_observables=True) + return dets.astype(np.float64), obs.reshape(-1).astype(bool) + except TypeError: + result = sampler.sample(shots, bit_packed=False) + + if isinstance(result, tuple): + if len(result) >= 2: + return result[0].astype( + np.float64), result[1].reshape(-1).astype(bool) + raise RuntimeError( + "Could not sample observables from DEM sampler for this Stim version. " + "Use generated-circuit mode or pass --diagnostics-only.") + + +def sample_problem(problem: Problem, shots: int, seed: int): + if problem.circuit is not None: + return sample_from_circuit(problem.circuit, shots, seed) + return sample_from_dem(problem.dem, shots, seed) + + +def path_largest_intermediate(info: Any) -> float: + value = getattr(info, "largest_intermediate", None) + if value is None: + return float("nan") + try: + return float(value) + except (TypeError, ValueError, OverflowError): + return float("nan") + + +def path_opt_cost(info: Any) -> float: + value = getattr(info, "opt_cost", None) + if value is None: + return float("nan") + try: + return float(value) + except (TypeError, ValueError, OverflowError): + return float("nan") + + +def einsum_rank(eq: str) -> int: + if "->" not in eq: + return 0 + lhs, rhs = eq.split("->", 1) + terms = [term for term in lhs.split(",") if term] + terms.append(rhs) + return max((len(term) for term in terms), default=0) + + +def path_max_einsum_rank(info: Any) -> int: + ranks = [ + einsum_rank(step[2]) for step in getattr(info, "contraction_list", ()) + ] + return max(ranks, default=0) + + +def fmt_float(value: float) -> str: + if math.isnan(value): + return "n/a" + return f"{value:.6e}" + + +def element_size(dtype: str) -> int: + return torch.tensor([], dtype=getattr(torch, dtype)).element_size() + + +def clear_cuda(device: str): + gc.collect() + if "cuda" in device and torch.cuda.is_available(): + torch.cuda.synchronize(torch.device(device)) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(torch.device(device)) + + +def peak_mem_mib(device: str) -> str: + if "cuda" not in device or not torch.cuda.is_available(): + return "n/a" + torch.cuda.synchronize(torch.device(device)) + return f"{torch.cuda.max_memory_allocated(torch.device(device)) / 2**20:.1f}" + + +def make_initial_priors(priors: np.ndarray, mode: str) -> np.ndarray: + if mode == "true": + return priors.copy() + if mode == "uniform": + return np.full_like(priors, priors.mean()) + raise ValueError(f"unknown init mode: {mode}") + + +def probs_to_logits(probs: np.ndarray, dtype: str) -> torch.Tensor: + eps = 1e-6 if dtype == "float32" else 1e-12 + clipped = np.clip(probs, eps, 1.0 - eps) + return torch.as_tensor(np.log(clipped / (1.0 - clipped)), + dtype=getattr(torch, dtype)) + + +def construct_optimizer(problem: Problem, syn: np.ndarray, flips: np.ndarray, + priors: np.ndarray, args: argparse.Namespace, + execute: str) -> NMOptimizer: + return NMOptimizer(problem.H, + problem.logical, + priors.tolist(), + syn, + flips, + dtype=args.dtype, + device=args.device, + execute=execute, + compile=args.compile) + + +def run_case(problem: Problem, + batch_size: int, + execute: str, + args: argparse.Namespace, + fixed_path: Any | None = None, + fixed_path_source: str = "") -> tuple[dict[str, str], Any | None]: + label = f"batch={batch_size}, execute={execute}" + row = { + "batch": str(batch_size), + "construct_batch": str(batch_size), + "execute": execute, + "path_source": fixed_path_source if fixed_path is not None else "self", + "status": "FAIL", + "construct_s": "n/a", + "largest_intermediate": "n/a", + "intermediate_gib": "n/a", + "opt_cost": "n/a", + "max_einsum_rank": "n/a", + "step_s": "n/a", + "peak_mem_mib": "n/a", + "error": "", + } + init_priors = make_initial_priors(problem.priors, args.init) + syn, flips = sample_problem(problem, batch_size, args.seed + batch_size) + + selected_path = None + try: + clear_cuda(args.device) + with timed(f"Construct NMOptimizer ({label})"): + t0 = time.perf_counter() + if args.run_step: + logits = probs_to_logits(init_priors, args.dtype) + logits = logits.to(torch.device(args.device)) + logits.requires_grad_(True) + torch_opt = torch.optim.Adam([logits], lr=args.lr) + progress_every = max(1, args.progress_every) + + def _progress(done: int, total: int, elapsed: float) -> None: + if done == 1 or done == total or done % progress_every == 0: + rate = done / elapsed if elapsed > 0 else float("nan") + print( + f" batch slice {done}/{total}: " + f"elapsed={elapsed:.1f}s, rate={rate:.3f}/s, " + f"peak_mem_mib={peak_mem_mib(args.device)}", + flush=True) + + opt, step_fn = make_training_step( + problem.H, + problem.logical, + init_priors.tolist(), + syn, + flips, + logits, + torch_opt, + batch_slicing="auto", + batch_slice_size=args.batch_slice_size, + progress_callback=_progress, + dtype=args.dtype, + device=args.device, + execute=execute, + compile=args.compile) + else: + opt = construct_optimizer(problem, syn, flips, init_priors, + args, execute) + step_fn = None + if fixed_path is not None: + print(f"Reusing reduced path from {fixed_path_source}", + flush=True) + opt.optimize_path(fixed_path) + row["construct_s"] = f"{time.perf_counter() - t0:.3f}" + + construct_batch = int(getattr(opt, "_batch_size", batch_size)) + row["construct_batch"] = str(construct_batch) + run_label = label + if construct_batch != batch_size: + run_label += f", construct_batch={construct_batch}" + + info = getattr(opt, "_reduced_info", None) + selected_path = getattr(opt, "path_batch", None) + largest = path_largest_intermediate(info) + cost = path_opt_cost(info) + row["largest_intermediate"] = fmt_float(largest) + row["opt_cost"] = fmt_float(cost) + row["max_einsum_rank"] = str(path_max_einsum_rank(info)) + if not math.isnan(largest): + gib = largest * element_size(args.dtype) / 2**30 + row["intermediate_gib"] = f"{gib:.3f}" + print( + "Path info: largest_intermediate=" + f"{row['largest_intermediate']}, estimated_one_tensor_GiB=" + f"{row['intermediate_gib']}, opt_cost={row['opt_cost']}", + flush=True) + + if args.run_step: + with timed(f"Adam step ({run_label})"): + t0 = time.perf_counter() + loss = step_fn() + if "cuda" in args.device and torch.cuda.is_available(): + torch.cuda.synchronize(torch.device(args.device)) + row["step_s"] = f"{time.perf_counter() - t0:.3f}" + print(f"Loss={float(loss.detach().cpu()):.6e}", flush=True) + + row["peak_mem_mib"] = peak_mem_mib(args.device) + row["status"] = "OK" + return row, selected_path + except Exception as exc: # intentionally broad: this is an OOM probe. + row["error"] = repr(exc) + with contextlib.suppress(Exception): + row["peak_mem_mib"] = peak_mem_mib(args.device) + print(f"FAILED CASE {label}: {exc!r}", flush=True) + if args.stop_on_failure: + raise + return row, selected_path + finally: + with contextlib.suppress(Exception): + del opt # type: ignore[name-defined] + gc.collect() + if "cuda" in args.device and torch.cuda.is_available(): + with contextlib.suppress(Exception): + torch.cuda.empty_cache() + + +def parse_execute(value: str) -> list[str]: + if value == "all": + return list(EXECUTE_MODES) + modes = [v.strip() for v in value.split(",") if v.strip()] + bad = sorted(set(modes) - set(EXECUTE_MODES)) + if bad: + raise ValueError(f"invalid execute mode(s): {bad}") + return modes + + +def make_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dem-path", + type=str, + default=None, + help="Optional exact .dem file to benchmark.") + parser.add_argument("--code", default="surface_code:rotated_memory_z") + parser.add_argument("--distance", type=int, default=5) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument( + "--noise-p", + type=float, + default=0.005, + help="Default probability for all enabled Stim noise mechanisms.") + parser.add_argument("--after-clifford-p", type=float, default=None) + parser.add_argument("--after-reset-p", type=float, default=None) + parser.add_argument("--before-measure-p", type=float, default=None) + parser.add_argument("--before-round-data-p", type=float, default=None) + parser.add_argument("--batch-sizes", + type=int, + nargs="+", + default=[1, 2, 4, 1000]) + parser.add_argument( + "--execute", + default="all", + help="all, codegen, unrolled, opt_einsum, cutensornet, or comma list") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", + choices=["float32", "float64"], + default="float64") + parser.add_argument("--init", + choices=["uniform", "true"], + default="uniform") + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--run-step", + action="store_true", + help="Actually execute one forward/backward Adam step. May OOM.") + parser.add_argument( + "--batch-slice-size", + type=int, + default=None, + help=("Construct with this smaller batch size and stream the " + "requested batch through that fixed path.")) + parser.add_argument("--microbatch-size", + dest="batch_slice_size", + type=int, + help=argparse.SUPPRESS) + parser.add_argument("--progress-every", + type=int, + default=10, + help="Print batch-slice progress every N chunks.") + parser.add_argument( + "--reuse-first-path", + action="store_true", + help=("For each requested batch size, reuse the first execute mode's " + "reduced contraction path for later execute modes.")) + parser.add_argument("--lr", + type=float, + default=1e-2, + help="Adam learning rate used with --run-step.") + parser.add_argument("--compile", + action="store_true", + help="Pass compile=True to NMOptimizer.") + parser.add_argument("--stop-on-failure", action="store_true") + return parser + + +def main() -> int: + parser = make_parser() + args = parser.parse_args() + if args.after_clifford_p is None: + args.after_clifford_p = args.noise_p + if args.after_reset_p is None: + args.after_reset_p = args.noise_p + if args.before_measure_p is None: + args.before_measure_p = args.noise_p + if args.before_round_data_p is None: + args.before_round_data_p = args.noise_p + + if args.batch_slice_size is not None and args.batch_slice_size <= 0: + parser.error("--batch-slice-size must be a positive integer.") + + execute_modes = parse_execute(args.execute) + print( + "Config: " + f"source={'dem-file' if args.dem_path else 'stim-generated'}, " + f"d={args.distance}, r={args.rounds}, device={args.device}, " + f"dtype={args.dtype}, execute={execute_modes}, " + f"batch_sizes={args.batch_sizes}, run_step={args.run_step}, " + f"reuse_first_path={args.reuse_first_path}, " + f"batch_slice_size={args.batch_slice_size}", + flush=True) + + with timed("Build/load DEM"): + problem = build_problem(args) + print( + f"DEM shape: H={problem.H.shape}, logical={problem.logical.shape}, " + f"priors={problem.priors.shape}", + flush=True) + print( + f"True priors: mean={problem.priors.mean():.6e}, " + f"min={problem.priors.min():.6e}, max={problem.priors.max():.6e}", + flush=True) + if problem.H.shape != (120, 1677): + print( + "WARNING: DEM shape is not Nico's reported d=5/r=5 size " + "H=(120, 1677). If you have the exact DEM, rerun with --dem-path.", + flush=True) + + rows = [] + for batch_size in args.batch_sizes: + fixed_path = None + fixed_path_source = "" + for execute in execute_modes: + row, selected_path = run_case(problem, + batch_size, + execute, + args, + fixed_path=fixed_path, + fixed_path_source=fixed_path_source) + rows.append(row) + if (args.reuse_first_path and fixed_path is None and + row["status"] == "OK" and selected_path is not None): + fixed_path = selected_path + fixed_path_source = f"execute={execute}" + + print("\n## NMOptimizer d=5/r=5 Large-DEM Memory Matrix") + headers = [ + "batch", "construct_batch", "execute", "path_source", "status", + "construct_s", "largest_intermediate", "intermediate_gib", "opt_cost", + "max_einsum_rank", "step_s", "peak_mem_mib", "error" + ] + print("| " + " | ".join(headers) + " |") + print("| " + " | ".join(["---"] * len(headers)) + " |") + for row in rows: + print("| " + " | ".join(row[h].replace("\n", " ") for h in headers) + + " |") + + return 0 if all(r["status"] == "OK" for r in rows) else 1 + + +if __name__ == "__main__": + raise SystemExit(main())