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/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..16d9c5146 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 @@ -20,6 +20,7 @@ import warnings from typing import Any, 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, ) @@ -106,6 +109,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 +156,65 @@ 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 _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.") + + _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. @@ -199,6 +267,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):: @@ -240,42 +311,47 @@ 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" + + # 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 +361,21 @@ 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() + + self._set_contractor("oe_torch_compiled", 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 +384,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 +526,277 @@ 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_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, + } + 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_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 +804,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. @@ -599,10 +846,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 +855,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 +901,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 +943,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 +1075,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 +1108,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 +1162,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 +1171,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, diff --git a/libs/qec/python/tests/test_nm_optimizer.py b/libs/qec/python/tests/test_nm_optimizer.py index 97dceaba1..1e9cfa06a 100644 --- a/libs/qec/python/tests/test_nm_optimizer.py +++ b/libs/qec/python/tests/test_nm_optimizer.py @@ -111,6 +111,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 +149,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 @@ -329,6 +357,74 @@ 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) + + # -- numerical guards -------------------------------------------------------- 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():