diff --git a/python/tests/beagle.py b/python/tests/beagle.py
index 8f71a68074..207957d8b2 100644
--- a/python/tests/beagle.py
+++ b/python/tests/beagle.py
@@ -314,20 +314,23 @@ def _compute_state_probability_matrix(fm, bm, ref_h, query_h, rho, mu):
         backward haplotype probability matrix.
     :rtype: tuple
     """
-    m = ref_h.shape[0]
-    h = ref_h.shape[1]
+    m = fm.shape[0]
+    h = fm.shape[1]
+    assert (m, h) == bm.shape
+    assert not np.any(fm < 0), "Forward probability matrix has negative values."
+    assert not np.any(np.isnan(fm)), "Forward probability matrix has NaN values."
+    assert not np.any(bm < 0), "Backward probability matrix has negative values."
+    assert not np.any(np.isnan(bm)), "Backward probability matrix has NaN values."
+    assert np.all(
+        np.isin(ref_h, [0, 1])
+    ), "Reference haplotypes have non-biallelic values."
+    assert np.all(np.isin(query_h, [0, 1])), "Query haplotype has non-biallelic values."
     assert m == len(query_h)
     assert m == len(rho)
     assert m == len(mu)
     assert 0 <= np.min(rho) and np.max(rho) <= 1
-    # BEAGLE caps mismatch probabilities at 0.5.
     assert 0 <= np.min(mu) and np.max(mu) <= 0.5
-    assert fm.shape == (m, h)
-    assert bm.shape == (m, h)
-    # Check all biallelic sites
-    assert np.all(np.isin(np.unique(ref_h), [0, 1]))
-    assert np.all(np.isin(np.unique(query_h), [-1, 0, 1]))
-    sm = np.zeros((m, h), dtype=np.float64)  # HMM state probability matrix
+    sm = np.zeros((m, h), dtype=np.float64)
     fwd_hap_probs = np.zeros((m, 4), dtype=np.float64)
     bwd_hap_probs = np.zeros((m, 4), dtype=np.float64)
     for i in np.arange(m - 1, -1, -1):