diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
index 6c6fffe74d..55320f451b 100644
--- a/python/tests/beagle_numba.py
+++ b/python/tests/beagle_numba.py
@@ -128,9 +128,12 @@ def get_alleles_at_site(self, i):
         return a1, ap1, a2, ap2
 
 
+""" Helper functions. """
+
+
 def remap_alleles(a):
     """
-    Helper function to map an allele encoded as a character to an integer.
+    Map an allele encoded as a character to an integer.
 
     TODO: Move elsewhere.
 
@@ -213,6 +216,9 @@ def convert_to_genetic_map_positions(pos, genetic_map=None):
     return est_cm
 
 
+""" Li & Stephens HMM. """
+
+
 def get_mismatch_probs(pos, error_rate):
     """
     Compute mismatch probabilities at genotyped positions.
@@ -363,6 +369,9 @@ def compute_state_prob_matrix(fm, bm):
     return sm
 
 
+""" Genotype imputation. """
+
+
 @njit
 def get_weights(typed_pos, untyped_pos, typed_cm, untyped_cm):
     """
@@ -640,6 +649,9 @@ def run_tsimpute(
     return (imputed_alleles, max_allele_probs)
 
 
+""" Evaluation metrics and printing of results. """
+
+
 # Individual-level data.
 def compute_individual_scores(alleles_1, allele_probs_1, alleles_2, allele_probs_2):
     """