diff --git a/python/tests/test_beagle.py b/python/tests/test_beagle.py
index 723d3e3760..39b25345d0 100644
--- a/python/tests/test_beagle.py
+++ b/python/tests/test_beagle.py
@@ -325,26 +325,27 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h, rho, mu):
     return (sm, fwd_hap_probs, bwd_hap_probs)
 
 
-def interpolate_haplotype_probability_matrix(
+def interpolate_allele_probabilities(
     fwd_hap_probs, bwd_hap_probs, genotyped_pos, imputed_pos
 ):
     """
-    Compute the interpolated haplotype probability matrix of a query haplotype.
+    Compute the interpolated allele probabilities at the imputed markers
+    of a query haplotype.
 
     The forward and backward haplotype probability matrices are of size (m, 2),
-    where m is the number of genotyped markers in the query haplotype.
+    where m is the number of genotyped markers.
     Here, we assume all biallelic sites, hence the 2.
 
     The interpolated haplotype probability matrix is of size (x, 2),
-    where x is the number of imputed markers in the query haplotype.
+    where x is the number of imputed markers.
     Again, we assume all biallelic sites, hence the 2.
 
     This implementation is based on Equation 2 (the rearranged one) in BB2016.
 
-    :param fwd_hap_probs: Forward haplotype probability matrix.
-    :param bwd_hap_probs: Backward haplotype probability matrix.
-    :param genotyped_pos: Site positions of genotyped markers.
-    :param imputed_pos: Site positions of imputed markers.
+    :param numpy.ndarray fwd_hap_probs: Forward haplotype probability matrix.
+    :param numpy.ndarray bwd_hap_probs: Backward haplotype probability matrix.
+    :param numpy.ndarray genotyped_pos: Site positions of genotyped markers.
+    :param numpy.ndarray imputed_pos: Site positions of imputed markers.
     :return: Interpolated haplotype probability matrix.
     :rtype: numpy.ndarray
     """
@@ -535,7 +536,7 @@ def run_beagle(ref_h, query_h, pos):
     assert bwd_hap_probs.shape == (m, 2)
     # Interpolate haplotype probabilities
     # from genotype markers to imputed markers
-    i_hap_probs = interpolate_haplotype_probability_matrix(
+    i_hap_probs = interpolate_allele_probabilities(
         fwd_hap_probs, bwd_hap_probs, genotyped_pos, imputed_pos
     )
     assert i_hap_probs.shape == (x, 2)