def _run_compare_with_baseline_sweep(self, batch_by_size_impl):
        """Compare reference batch_by_size implementation with batch_by_size_baseline
        across a dense grid of hyperparam values"""
        MAX_MAX_TOKENS = 10
        NUM_TOKENS_VECS_COUNT = 5
        for indices_len in [10, 11]:  # try odd and even len of indices
            for max_sentences in range(0, indices_len + 2):
                for max_tokens in range(0, MAX_MAX_TOKENS):
                    for bsz_mult in range(1, max(MAX_MAX_TOKENS, indices_len) + 2):
                        for _ in range(NUM_TOKENS_VECS_COUNT):
                            num_tokens_vec = np.random.randint(
                                0, max_tokens + 1, size=indices_len
                            )
                            self._compare_results(
                                indices_len,
                                batch_by_size_impl,
                                max_sentences,
                                max_tokens,
                                bsz_mult,
                                num_tokens_vec,
                            )
Compare reference batch_by_size implementation with batch_by_size_baseline
        across a dense grid of hyperparam values
edgelm/tests/test_data_utils.py