diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
index fea4f5f00f..15db7ce786 100644
--- a/python/tests/beagle_numba.py
+++ b/python/tests/beagle_numba.py
@@ -338,7 +338,7 @@ def interpolate_allele_prob(sm, ref_h, typed_pos, untyped_pos, typed_cm, untyped
     return probs_rescaled
 
 
-def get_map_alleles(allele_prob):
+def get_map_alleles(allele_probs):
     """
     Compute the maximum a posteriori alleles at ungenotyped markers of a query haplotype,
     based on posterior marginal allele probabilities.
@@ -346,15 +346,15 @@ def get_map_alleles(allele_prob):
     The imputed alleles is an array of size x.
 
     WARN: If the allele probabilities are equal, then allele 0 is arbitrarily chosen.
-    TODO: Analyse how often this happens and the effect of this arbitrary choice.
+    TODO: Investigate how often this happens and the effect of this arbitrary choice.
 
-    :param numpy.ndarray allele_prob: Interpolated allele probabilities.
+    :param numpy.ndarray allele_probs: Interpolated allele probabilities.
     :return: Imputed alleles and their associated probabilities.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """
-    imputed_alleles = np.argmax(allele_prob, axis=1)
-    max_allele_prob = np.max(allele_prob, axis=1)
-    return (imputed_alleles, max_allele_prob)
+    imputed_alleles = np.argmax(allele_probs, axis=1)
+    max_allele_probs = np.max(allele_probs, axis=1)
+    return (imputed_alleles, max_allele_probs)
 
 
 def run_interpolation_beagle(