From 9de9dbb44160fc149be811b55c21ce9bc07d7e61 Mon Sep 17 00:00:00 2001 From: vedika-saravanan Date: Tue, 19 May 2026 17:12:11 -0400 Subject: [PATCH 1/3] docs(qec): document tensor network noise-learning decoders Adds Sphinx API narrative and decoder docs for noise-learning integration, split from the integration PR for focused review. Signed-off-by: vedika-saravanan --- .../api/qec/tensor_network_decoder_api.rst | 183 +++++++++++++++++- docs/sphinx/components/qec/introduction.rst | 21 ++ docs/sphinx/examples_rst/qec/decoders.rst | 25 +++ 3 files changed, 228 insertions(+), 1 deletion(-) diff --git a/docs/sphinx/api/qec/tensor_network_decoder_api.rst b/docs/sphinx/api/qec/tensor_network_decoder_api.rst index b5bbc600c..16279de2e 100644 --- a/docs/sphinx/api/qec/tensor_network_decoder_api.rst +++ b/docs/sphinx/api/qec/tensor_network_decoder_api.rst @@ -93,4 +93,185 @@ :param optimize: Optimization options or None :param batch_size: (int, optional) Batch size for optimization (default: -1, no batching) - :returns: Optimizer info object \ No newline at end of file + :returns: Optimizer info object + +.. class:: cudaq_qec.plugins.decoders.tensor_network_decoder.NMOptimizer + + Differentiable noise-model optimizer built on top of :class:`TensorNetworkDecoder`. + + Fits a factorised per-error noise model to a syndrome dataset by + backpropagating through a torch-backed tensor-network contraction. + The noise probabilities are maintained as ``torch`` tensors with + ``requires_grad=True`` so they can be updated with any ``torch.optim`` + optimizer. + + Requires Python 3.11 or higher and the same optional dependencies as + :class:`TensorNetworkDecoder` (``pip install cudaq-qec[tensor-network-decoder]``). + PyTorch must also be installed. + + .. note:: + Quick-start example (logit-space training; the loss has no ``log`` + guard, so direct probability training requires per-step clamping + into ``[eps, 1 - eps]``):: + + import numpy as np + import torch + from cudaq_qec.plugins.decoders.tensor_network_decoder import ( + NMOptimizer, make_compiled_step, + ) + + H = np.array([[1, 1, 0], [0, 1, 1]], dtype=np.float64) + logical = np.array([[1, 0, 1]], dtype=np.float64) + priors = [0.1, 0.2, 0.3] + + opt = NMOptimizer(H, logical, priors, syndrome_data, obs_flips, + dtype="float64") + logits = torch.logit(opt.noise_params[0].detach()).requires_grad_() + adam = torch.optim.Adam([logits], lr=0.01) + step = make_compiled_step(opt, logits, adam) + for _ in range(100): + step() + + :param H: Parity check matrix (numpy.ndarray), shape (num_checks, num_errors) + :param logical_obs: Logical observable matrix (numpy.ndarray), shape (1, num_errors) + :param noise_model: Initial per-error probabilities, list of floats in (0, 1). + Values outside ``[eps, 1 - eps]`` are clamped at + construction with a ``UserWarning``; non-finite values + raise ``ValueError``. ``eps`` is ``1e-12`` for + ``"float64"`` and ``1e-6`` for ``"float32"``. + :param syndrome_data: Observed syndromes, numpy.ndarray of shape (num_shots, num_checks) + :param observable_flips: Observed logical flips, bool array of length num_shots + :param check_inds: (optional) List of check index names; defaults track the parent decoder. + :param error_inds: (optional) List of error index names; defaults track the parent decoder. + :param logical_inds: (optional) List of logical index names; defaults track the parent decoder. + :param logical_tags: (optional) List of logical tags; defaults track the parent decoder. + :param dtype: (str, optional) ``"float32"`` (default) or ``"float64"``; + other values raise ``ValueError``. + :param device: (str, optional) Torch device, e.g. ``"cpu"`` or ``"cuda"`` (default: ``"cuda"``) + :param compile: (bool, optional, keyword-only) If ``True``, wrap the forward + and loss in :func:`torch.compile`. Most useful with + ``execute="codegen"``. Defaults to ``False``. + :param execute: (str, optional, keyword-only) Forward backend. ``"codegen"`` + (default) partial-evaluates the contraction path into a flat + Python function with named locals; ``"unrolled"`` keeps an + interpretive einsum list; ``"opt_einsum"`` dispatches via + :func:`opt_einsum.contract_expression`. + :param compile_mode: (str, optional, keyword-only) Forwarded to + :func:`torch.compile` (e.g. ``"reduce-overhead"``, + ``"default"``); ignored when ``compile=False``. + :param dynamic_syndromes: (bool, optional, keyword-only) If ``True`` + (default), syndromes are runtime arguments to the + compiled forward, so :meth:`update_dataset` reuses + the codegen/``torch.compile`` artifact when shapes + are unchanged. ``False`` bakes syndromes into the + closure -- faster per call but every + :meth:`update_dataset` rebuilds the graph. Only + affects ``execute="codegen"``. + + **Attributes** + + .. attribute:: noise_params + + ``list[torch.Tensor]`` — the learnable noise-probability tensors; pass + directly to a ``torch.optim`` optimizer. + + .. attribute:: torch_device + + ``torch.device`` derived from the ``device`` constructor argument. + Read-only. + + .. attribute:: observable_flips + + Bool ``torch.Tensor`` of logical flip outcomes for the current + syndrome batch. Assigning a new value also rebuilds the fused + loss closure (the observable indices are baked into the codegen); + prefer :meth:`update_dataset` when swapping syndromes and flips + together. + + **Methods** + + .. method:: current_syndrome_args() + + Return the syndrome argument expected by the callable from + :meth:`loss_fn`: the live tuple when ``dynamic_syndromes=True``, + or ``()`` for static codegen (syndromes are closure-baked). + Re-fetch each step so an intervening :meth:`update_dataset` is + reflected. + + :returns: ``tuple[torch.Tensor, ...]`` + + .. method:: cross_entropy_loss() + + Compute the cross-entropy loss between the predicted logical-flip + probabilities and the observed ``observable_flips``. + + :returns: Scalar ``torch.Tensor`` (differentiable). + + .. method:: decoder_prediction() + + Run the forward pass and return per-shot probabilities. + + :returns: ``torch.Tensor`` of shape ``(num_shots, 2)`` where column 1 + is ``P(logical flip | syndrome)``. + + .. method:: logical_error_rate() + + Fraction of shots where ``argmax`` of :meth:`decoder_prediction` + disagrees with :attr:`observable_flips`. Not differentiable + (runs under :func:`torch.no_grad`). + + :returns: ``float`` in ``[0, 1]``. + + .. method:: loss_fn(from_logits=True) + + Return a compiled callable ``fn(params, syndrome_tuple) -> loss`` + suitable for use with external optimizers or ``torch.compile``. + + :param from_logits: If ``True`` (default), ``params`` are interpreted + as logits and passed through ``sigmoid`` before + contraction. If ``False``, ``params`` are + interpreted as probabilities already in ``[0, 1]``. + :returns: Compiled loss function. + + .. method:: optimize_path(optimize=None, batch_size=-1) + + Cache a contraction path via quimb / opt_einsum and rebuild the + compiled forward. Pass e.g. ``cotengra.HyperOptimizer()`` to run a + more expensive path search; ``None`` falls back to ``"auto"``. + + :param optimize: Optimization options (e.g. a ``cotengra.HyperOptimizer``) + or ``None``. + :param batch_size: Accepted for signature compatibility; ignored. + :returns: Contraction info object. + + .. method:: update_dataset(syndrome_data, observable_flips, enforce_shape=True) + + Swap in a new syndrome batch without rebuilding the tensor network. + If ``dynamic_syndromes=True`` and the batch size is unchanged, the + compiled contraction path is reused; a shape change triggers a full + rebuild. + + :param syndrome_data: numpy.ndarray of shape (num_shots, num_checks) + :param observable_flips: bool array of length num_shots + :param enforce_shape: (bool, optional, default ``True``) Assert + per-tensor shapes match the existing layout + before patching in place. A batch-size change + triggers a full rebuild regardless. + +.. function:: cudaq_qec.plugins.decoders.tensor_network_decoder.make_compiled_step(optimizer, logits, torch_optimizer) + + Build a no-arg callable that runs one Adam step and returns the loss. + + The returned ``step()`` callable zeros gradients, evaluates the + optimizer's fused ``loss_fn(from_logits=True)`` (sigmoid + contraction + + cross-entropy), backpropagates, and steps ``torch_optimizer``. Intended + for training in logit space; pair with :class:`NMOptimizer` constructed + with ``compile=True`` for a ``torch.compile``-d variant. + + :param optimizer: An :class:`NMOptimizer` instance providing the fused + inner loss. + :param logits: Trainable 1-D ``torch.Tensor`` of length + ``len(optimizer.error_inds)`` with ``requires_grad=True``. + :param torch_optimizer: A ``torch.optim`` instance owning ``logits``. + :returns: A no-arg callable that performs one optimization step and + returns the scalar loss as a ``torch.Tensor``. \ No newline at end of file diff --git a/docs/sphinx/components/qec/introduction.rst b/docs/sphinx/components/qec/introduction.rst index 99756c73c..4088795a9 100644 --- a/docs/sphinx/components/qec/introduction.rst +++ b/docs/sphinx/components/qec/introduction.rst @@ -899,6 +899,27 @@ The decoder returns the probability that the logical observable has flipped for that this GPU will not be supported by the Tensor Network Decoder when CUDA-Q 0.5.0 is released. +Learning the Noise Model from Data +"""""""""""""""""""""""""""""""""" + +When the true per-error noise rates are unknown (typical of real hardware), +the Tensor Network Decoder ships with ``NMOptimizer``, a differentiable +extension that **fits the noise model directly from observed syndromes and +logical-flip outcomes**. Noise probabilities are held as PyTorch tensors +with ``requires_grad=True``; backpropagating through the tensor-network +contraction yields gradients that any ``torch.optim`` optimizer (Adam, SGD, +etc.) can update. Starting from a uniform initial prior and a few hundred +Adam steps is usually enough to recover the per-error rates and beat a +static-uniform baseline on a held-out batch. + +This is offline -- training happens once on a representative syndrome +dataset, and the learned probabilities can then be used as a standard +static noise model for batch decoding. See +:ref:`tensor_network_decoder_api_python` for the ``NMOptimizer`` API and +the *Learning Noise Models with NMOptimizer* example in +:doc:`../../examples_rst/qec/decoders` for a runnable end-to-end demo on a +Stim repetition-code circuit. + Sliding Window Decoder ^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/sphinx/examples_rst/qec/decoders.rst b/docs/sphinx/examples_rst/qec/decoders.rst index fe240ed1c..c86f0a516 100644 --- a/docs/sphinx/examples_rst/qec/decoders.rst +++ b/docs/sphinx/examples_rst/qec/decoders.rst @@ -136,6 +136,31 @@ See Also: - ``cudaq_qec.plugins.decoders.tensor_network_decoder`` +Learning Noise Models with NMOptimizer ++++++++++++++++++++++++++++++++++++++++ + +:class:`~cudaq_qec.plugins.decoders.tensor_network_decoder.NMOptimizer` extends +the Tensor Network Decoder with differentiable noise probabilities. Given a +batch of observed syndromes and logical-flip outcomes, it fits per-error noise +rates by backpropagating through the tensor-network contraction using PyTorch. + +The following example builds a distance-3 repetition-code circuit with +**asymmetric** noise (data-qubit depolarization is 10x measurement-flip +probability), samples syndromes from Stim, and trains +:class:`NMOptimizer` from a uniform initial prior with 300 Adam steps in +logit space. It then compares the **logical error rate (LER)** of the +learned noise model against a static uniform-prior baseline on a 20k-shot +held-out batch — demonstrating that fitting per-error rates from data +decodes meaningfully better than assuming uniform noise: + +.. literalinclude:: ../../examples/qec/python/noise_learning.py + :language: python + :start-after: [Begin Documentation] + +See Also: + +- :ref:`tensor_network_decoder_api_python` + .. _deploying-ai-decoders: Deploying AI Decoders with TensorRT From d200171266a1598cbfb7fc1e251a949f5c7d0998 Mon Sep 17 00:00:00 2001 From: vedika-saravanan Date: Mon, 22 Jun 2026 17:16:09 -0400 Subject: [PATCH 2/3] document tensor network noise-learning decoders Signed-off-by: vedika-saravanan --- .../api/qec/tensor_network_decoder_api.rst | 22 +++++++++++-------- docs/sphinx/components/qec/introduction.rst | 6 ++--- docs/sphinx/examples_rst/qec/decoders.rst | 6 ++--- 3 files changed, 19 insertions(+), 15 deletions(-) diff --git a/docs/sphinx/api/qec/tensor_network_decoder_api.rst b/docs/sphinx/api/qec/tensor_network_decoder_api.rst index 16279de2e..bbc1eef38 100644 --- a/docs/sphinx/api/qec/tensor_network_decoder_api.rst +++ b/docs/sphinx/api/qec/tensor_network_decoder_api.rst @@ -95,7 +95,7 @@ :param batch_size: (int, optional) Batch size for optimization (default: -1, no batching) :returns: Optimizer info object -.. class:: cudaq_qec.plugins.decoders.tensor_network_decoder.NMOptimizer +.. class:: cudaq_qec.NMOptimizer Differentiable noise-model optimizer built on top of :class:`TensorNetworkDecoder`. @@ -110,15 +110,19 @@ PyTorch must also be installed. .. note:: - Quick-start example (logit-space training; the loss has no ``log`` - guard, so direct probability training requires per-step clamping - into ``[eps, 1 - eps]``):: + + ``NMOptimizer`` uses torch-backed contractions so gradients can flow + through the tensor-network contraction. cuTensorNet acceleration is + supported by :class:`TensorNetworkDecoder` for decoding, but is not + used by ``NMOptimizer``. + + .. note:: + Quick-start example (logit-space training; direct probability + training can saturate if an optimizer step leaves ``[0, 1]``):: import numpy as np import torch - from cudaq_qec.plugins.decoders.tensor_network_decoder import ( - NMOptimizer, make_compiled_step, - ) + from cudaq_qec import NMOptimizer, make_compiled_step H = np.array([[1, 1, 0], [0, 1, 1]], dtype=np.float64) logical = np.array([[1, 0, 1]], dtype=np.float64) @@ -172,7 +176,7 @@ .. attribute:: noise_params - ``list[torch.Tensor]`` — the learnable noise-probability tensors; pass + ``list[torch.Tensor]`` -- the learnable noise-probability tensors; pass directly to a ``torch.optim`` optimizer. .. attribute:: torch_device @@ -258,7 +262,7 @@ before patching in place. A batch-size change triggers a full rebuild regardless. -.. function:: cudaq_qec.plugins.decoders.tensor_network_decoder.make_compiled_step(optimizer, logits, torch_optimizer) +.. function:: cudaq_qec.make_compiled_step(optimizer, logits, torch_optimizer) Build a no-arg callable that runs one Adam step and returns the loss. diff --git a/docs/sphinx/components/qec/introduction.rst b/docs/sphinx/components/qec/introduction.rst index f328cc19d..d75b4df62 100644 --- a/docs/sphinx/components/qec/introduction.rst +++ b/docs/sphinx/components/qec/introduction.rst @@ -912,9 +912,9 @@ extension that **fits the noise model directly from observed syndromes and logical-flip outcomes**. Noise probabilities are held as PyTorch tensors with ``requires_grad=True``; backpropagating through the tensor-network contraction yields gradients that any ``torch.optim`` optimizer (Adam, SGD, -etc.) can update. Starting from a uniform initial prior and a few hundred -Adam steps is usually enough to recover the per-error rates and beat a -static-uniform baseline on a held-out batch. +etc.) can update. Starting from a uniform initial prior, Adam optimization +can learn non-uniform per-error rates and improve over a static-uniform +baseline on a held-out batch. This is offline -- training happens once on a representative syndrome dataset, and the learned probabilities can then be used as a standard diff --git a/docs/sphinx/examples_rst/qec/decoders.rst b/docs/sphinx/examples_rst/qec/decoders.rst index dfaf19b7c..3bd12a828 100644 --- a/docs/sphinx/examples_rst/qec/decoders.rst +++ b/docs/sphinx/examples_rst/qec/decoders.rst @@ -153,7 +153,7 @@ See Also: Learning Noise Models with NMOptimizer +++++++++++++++++++++++++++++++++++++++ -:class:`~cudaq_qec.plugins.decoders.tensor_network_decoder.NMOptimizer` extends +:class:`~cudaq_qec.NMOptimizer` extends the Tensor Network Decoder with differentiable noise probabilities. Given a batch of observed syndromes and logical-flip outcomes, it fits per-error noise rates by backpropagating through the tensor-network contraction using PyTorch. @@ -164,10 +164,10 @@ probability), samples syndromes from Stim, and trains :class:`NMOptimizer` from a uniform initial prior with 300 Adam steps in logit space. It then compares the **logical error rate (LER)** of the learned noise model against a static uniform-prior baseline on a 20k-shot -held-out batch — demonstrating that fitting per-error rates from data +held-out batch -- demonstrating that fitting per-error rates from data decodes meaningfully better than assuming uniform noise: -.. literalinclude:: ../../examples/qec/python/noise_learning.py +.. literalinclude:: ../../examples/qec/python/tn_noise_learning.py :language: python :start-after: [Begin Documentation] From 00a42907663d28e8e01d355a34d57d32c740f5cd Mon Sep 17 00:00:00 2001 From: vedika-saravanan Date: Tue, 7 Jul 2026 10:48:45 -0400 Subject: [PATCH 3/3] fix syndrome_data and obs_flips in docs example Signed-off-by: vedika-saravanan --- docs/sphinx/api/qec/tensor_network_decoder_api.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/sphinx/api/qec/tensor_network_decoder_api.rst b/docs/sphinx/api/qec/tensor_network_decoder_api.rst index bbc1eef38..7cdb50a13 100644 --- a/docs/sphinx/api/qec/tensor_network_decoder_api.rst +++ b/docs/sphinx/api/qec/tensor_network_decoder_api.rst @@ -128,6 +128,10 @@ logical = np.array([[1, 0, 1]], dtype=np.float64) priors = [0.1, 0.2, 0.3] + rng = np.random.default_rng(0) + syndrome_data = rng.integers(0, 2, size=(200, 2)).astype(np.float64) + obs_flips = rng.integers(0, 2, size=200).astype(bool) + opt = NMOptimizer(H, logical, priors, syndrome_data, obs_flips, dtype="float64") logits = torch.logit(opt.noise_params[0].detach()).requires_grad_()