Enable torch-only NMOptimizer execution modes

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,24 +37,57 @@
 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)
     H = np.zeros(matrices.check_matrix.shape)
     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__":

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,13 +145,20 @@ 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"),
     )
     _allowed_backends: ClassVar[list[str]] = ["numpy", "torch"]
     _contractors: ClassVar[dict[str, Callable]] = {
         "numpy": contractor,
         "torch": contractor,
+        "oe_torch": oe_torch_contractor,
+        "oe_torch_compiled": oe_torch_compiled_contractor,
         "cutensornet": cutn_contractor,
     }