Fix TensorNetworkDecoder float32 default causing non-physical posteriors

Signed-off-by: vedika-saravanan <vsaravanan@nvidia.com>

Author: vedika-saravanan

docs/sphinx/api/qec/tensor_network_decoder_api.rst

@@ -68,7 +68,7 @@
     :param logical_inds: (optional) List of logical index names
     :param logical_tags: (optional) List of logical tags
     :param contract_noise_model: (bool, optional) Whether to contract the noise model at initialization (default: True)
-    :param dtype: (str, optional) Data type for tensors (default: "float32")
+    :param dtype: (str, optional) Data type for tensors (default: "float64")
     :param device: (str, optional) Device for tensor operations ("cpu", "cuda", or "cuda:X", default: "cuda")
 
     **Methods**

libs/qec/python/cudaq_qec/plugins/decoders/tensor_network_decoder.py

@@ -22,6 +22,23 @@
     tensor_network_from_syndrome_batch, tensor_network_from_logical_observable)
 
 
+def _safe_posterior(mass0, mass1) -> tuple:
+    """Validate contraction masses and compute the posterior P(logical flip).
+
+    Returns (posterior, is_valid). On any non-physical input (negative, NaN,
+    infinite, or zero-sum masses) returns (0.5, False) so callers can mark the
+    result unconverged rather than silently propagating a bad value.
+    """
+    m0, m1 = float(mass0), float(mass1)
+    if (np.isnan(m0) or np.isnan(m1) or np.isinf(m0) or np.isinf(m1) or
+            m0 < 0.0 or m1 < 0.0):
+        return 0.5, False
+    denom = m0 + m1
+    if not np.isfinite(denom) or denom == 0.0:
+        return 0.5, False
+    return m1 / denom, True
+
+
 @qec.decoder("tensor_network_decoder")
 class TensorNetworkDecoder:
     r"""A general class for tensor network decoders.
@@ -110,7 +127,7 @@ def __init__(
         logical_inds: list[str] | None = None,
         logical_tags: list[str] | None = None,
         contract_noise_model: bool = True,
-        dtype: str = "float32",
+        dtype: str = "float64",
         device: str = "cuda",
     ) -> None:
         """Initialize a sparse representation of a tensor network decoder for an arbitrary code
@@ -406,12 +423,11 @@ def decode(
             device_id=self.contractor_config.device_id,
         )
 
+        posterior, valid = _safe_posterior(contraction_value[0],
+                                           contraction_value[1])
         res = qec.DecoderResult()
-        res.converged = True
-        res.result = [
-            float(contraction_value[1] /
-                  (contraction_value[1] + contraction_value[0]))
-        ]
+        res.converged = valid
+        res.result = [posterior]
         return res
 
     def decode_batch(
@@ -463,18 +479,20 @@ def decode_batch(
         )
 
         probabilities = []
+        converged = []
         for r in range(syndrome_batch.shape[0]):
-            probabilities.append(
-                float(contraction_value[r, 1] /
-                      (contraction_value[r, 1] + contraction_value[r, 0])))
+            posterior, valid = _safe_posterior(contraction_value[r, 0],
+                                               contraction_value[r, 1])
+            probabilities.append(posterior)
+            converged.append(valid)
 
         # Python `decode_batch` override: construct a BatchDecoderResult
         # directly, bypassing the native decoder aggregation path. This is
         # the only sanctioned caller of the BatchDecoderResult constructor;
         # see its docstring for the supported construction surface.
         return qec.BatchDecoderResult(
             np.asarray(probabilities, dtype=np.float64).reshape((-1, 1)),
-            np.ones(syndrome_batch.shape[0], dtype=bool),
+            np.asarray(converged, dtype=bool),
         )
 
     def optimize_path(

libs/qec/python/tests/test_tensor_network_decoder.py

@@ -62,7 +62,7 @@ def test_decoder_init_and_attributes():
         assert decoder.contractor_config.contractor_name == "torch"
         assert decoder.contractor_config.backend == "torch"
         assert decoder.contractor_config.device == "cpu"
-    assert decoder._dtype == "float32"
+    assert decoder._dtype == "float64"
 
 
 def test_decoder_replace_logical_observable():
@@ -585,5 +585,98 @@ def test_invalid_combo():
         ContractorConfig("torch", "numpy", "cpu")
 
 
+# ---------------------------------------------------------------------------
+# Focused tests for _safe_posterior / mass-validation logic, using a mocked
+# contractor so these run without a GPU and without a real contraction.
+# ContractorConfig is a frozen dataclass; its contractor property reads from
+# the class-level _contractors dict, which is the correct mock surface.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_decoder(monkeypatch):
+    """Fixture that returns a factory for decoders with a mocked contractor.
+
+    Usage:
+        decoder = mock_decoder(np.array([m0, m1]))           # single decode
+        decoder = mock_decoder(np.array([[...],...]), rows=N) # batch decode
+    The factory patches ContractorConfig._contractors for the active backend
+    and pre-sets path_single/path_batch so no path optimisation is triggered.
+    """
+
+    def _make(mock_output, rows=None):
+        H, logical, noise = make_simple_code()
+        decoder = qec.get_decoder("tensor_network_decoder",
+                                  H,
+                                  logical_obs=logical,
+                                  noise_model=noise)
+        decoder.path_single = "auto"
+        decoder.path_batch = "auto"
+        if rows is not None:
+            decoder._batch_size = rows
+        name = decoder.contractor_config.contractor_name
+        monkeypatch.setitem(ContractorConfig._contractors, name,
+                            lambda *a, **kw: mock_output)
+        return decoder
+
+    return _make
+
+
+def test_default_dtype_is_float64():
+    H, logical, noise = make_simple_code()
+    decoder = qec.get_decoder("tensor_network_decoder",
+                              H,
+                              logical_obs=logical,
+                              noise_model=noise)
+    assert decoder._dtype == "float64"
+
+
+_BIG = np.finfo(np.float64).max
+
+
+@pytest.mark.parametrize(
+    "masses",
+    [
+        [-0.1, 1.1],  # negative mass
+        [0.0, 0.0],  # zero denominator
+        [float("nan"), 0.5],  # NaN
+        [float("inf"), 0.5],  # infinite mass
+        [_BIG, _BIG],  # finite masses whose sum overflows to inf
+    ])
+def test_decode_invalid_masses_marks_unconverged(mock_decoder, masses):
+    res = mock_decoder(np.array(masses)).decode([0.0, 0.0])
+    assert not res.converged
+    assert res.result[0] == pytest.approx(0.5)
+
+
+def test_decode_valid_masses_converged(mock_decoder):
+    res = mock_decoder(np.array([0.3, 0.7])).decode([0.0, 0.0])
+    assert res.converged
+    assert res.result[0] == pytest.approx(0.7)
+
+
+@pytest.mark.parametrize(
+    "bad_row",
+    [
+        [-0.1, 1.1],  # negative mass
+        [float("nan"), 0.5],  # NaN
+        [0.0, 0.0],  # zero denominator
+    ])
+def test_decode_batch_invalid_row_marks_unconverged(mock_decoder, bad_row):
+    mock = np.array([[0.3, 0.7], bad_row])
+    res = mock_decoder(mock, rows=2).decode_batch(np.zeros((2, 2)))
+    assert res.converged[0]
+    assert not res.converged[1]
+    assert res.result[0, 0] == pytest.approx(0.7)
+    assert res.result[1, 0] == pytest.approx(0.5)
+
+
+def test_decode_batch_all_valid_all_converged(mock_decoder):
+    mock = np.array([[0.2, 0.8], [0.6, 0.4], [0.1, 0.9]])
+    res = mock_decoder(mock, rows=3).decode_batch(np.zeros((3, 2)))
+    assert np.all(res.converged)
+    np.testing.assert_allclose(res.result[:, 0], [0.8, 0.4, 0.9])
+
+
 if __name__ == "__main__":
     pytest.main()