diff --git a/python/tests/beagle.py b/python/tests/beagle.py
index 6e9d3592a1..fa826cb48a 100644
--- a/python/tests/beagle.py
+++ b/python/tests/beagle.py
@@ -428,8 +428,8 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h):
 
 def interpolate_allele_probabilities(sm, ref_h, genotyped_pos, imputed_pos):
     """
-    Compute the interpolated allele probabilities at imputed markers
-    of a query haplotype following Equation 1 of BB2016.
+    Compute the interpolated allele probabilities at imputed markers of
+    a query haplotype following Equation 1 of BB2016.
 
     Assuming all biallelic sites, the interpolated allele probability matrix
     is of size (x, 2), where x is the number of imputed markers.
@@ -437,8 +437,8 @@ def interpolate_allele_probabilities(sm, ref_h, genotyped_pos, imputed_pos):
     This function takes the output of `compute_state_probability_matrix`.
 
     Note that this function takes:
-    1. HMM state probability matrix across genotyped markers (size of m).
-    2. reference haplotypes subsetted to imputed markers (size of x).
+    1. HMM state probability matrix across genotyped markers of size (m, h).
+    2. reference haplotypes subsetted to imputed markers of size (x, h).
 
     :param numpy.ndarray sm: HMM state probability matrix at genotyped markers.
     :param numpy.ndarray ref_h: Reference haplotypes subsetted to imputed markers.
@@ -447,15 +447,15 @@ def interpolate_allele_probabilities(sm, ref_h, genotyped_pos, imputed_pos):
     :return: Interpolated allele probabilities.
     :rtype: numpy.ndarray
     """
-    h = ref_h.shape[1]
     m = len(genotyped_pos)
     x = len(imputed_pos)
-    assert sm.shape == (m, h)
-    assert ref_h.shape == (x, h)
+    h = ref_h.shape[1]
+    assert (m, h) == sm.shape
+    assert (x, h) == ref_h.shape
     genotyped_cm = convert_to_genetic_map_position(genotyped_pos)
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
     weights = get_weights(genotyped_pos, imputed_pos)
-    assert len(weights) == x
+    assert x == len(weights)
     p = np.zeros((x, 2), dtype=np.float64)
 
     def _compute_allele_probabilities(a):