    def test_repeated_ngram_blocking_end_indices(self):
        """
        Ensures that the ngram blocking does not mess up when one sequence is shorter
        than another, which would result in repeated "end" symbols.
        """
        # We block unigrams, but 5 (the end symbol) is repeated and it does not mess
        # up the sequence's probability
        step_function = get_step_function(repeated_ngram_transition_probabilities_0)
        self.beam_search.beam_size = 2
        self.beam_search.constraints = [RepeatedNGramBlockingConstraint(ngram_size=1)]
        expected_top_k = np.array([[1, 3, 5, 5], [1, 2, 3, 5]])
        expected_log_probs = np.log(np.array([0.6 * 1e-9, 0.4 * 1e-9]))
        self._check_results(
            expected_top_k=expected_top_k,
            expected_log_probs=expected_log_probs,
            take_step=step_function,
        )
Ensures that the ngram blocking does not mess up when one sequence is shorter
than another, which would result in repeated "end" symbols.
tests/nn/beam_search_test.py