    def test_take_repeated_ngram_step(self):
        """
        Tests to ensure the top-k from the `repeated_ngram_transition_probabilities_0`
        transition matrix is expected. The transitions are:

            - p(1|start) = 1.0
            - p(2|1) = 0.4
            - p(3|1) = 0.6
            - p(end|1) = 1e-9
            - p(3|2) = 1.0
            - p(end|2) = 1e-9
            - p(1|3) = 1.0
            - p(end|3) = 1e-9

        The probabilities don't add up 1 because of the 1e-9 transitions to end. That doesn't
        really matter. Each state just needed some transition to the end probability with a very
        small probability to ensure it's possible to reach the end state from there and that it
        isn't selected by beam search without a constraint.

        Below is the beam search tracing for beam size 2. Any sequence below the
        line is not selected by beam search. The number that comes before the sequence
        is the probability of the sequence.

        Step 1
        1.0: [1]

        Step 2
        0.6: [1, 3]
        0.4: [1, 2]
        -----
        1e-9: [1, 2, end]

        Step 3
        0.6: [1, 3, 1]
        0.4: [1, 2, 3]
        -----
        0.6 * 1e-9: [1, 3, end]
        0.4 * 1e-9: [1, 2, end]

        Step 4
        0.4:  [1, 2, 3, 1]
        0.36: [1, 3, 1, 3]
        -----
        0.24:       [1, 3, 1, 2]
        0.6 * 1e-9: [1, 3, 1, end]
        0.4 * 1e-9: [1, 2, 3, end]

        Step 5
        0.36: [1, 3, 1, 3, 1]
        0.24: [1, 2, 3, 1, 3]
        -----
        0.16:        [1, 2, 3, 1, 2]
        0.4 * 1e-9:  [1, 2, 3, 1, end]
        0.36 * 1e-9: [1, 3, 1, 3, end]
        """
        step_function = get_step_function(repeated_ngram_transition_probabilities_0)
        self.beam_search.beam_size = 2
        self.beam_search.max_steps = 5
        expected_top_k = np.array([[1, 3, 1, 3, 1], [1, 2, 3, 1, 3]])
        expected_log_probs = np.log(np.array([0.36, 0.24]))
        self._check_results(
            expected_top_k=expected_top_k,
            expected_log_probs=expected_log_probs,
            take_step=step_function,
        )
Tests to ensure the top-k from the `repeated_ngram_transition_probabilities_0`
transition matrix is expected. The transitions are:

- p(1|start) = 1.0
- p(2|1) = 0.4
- p(3|1) = 0.6
- p(end|1) = 1e-9
- p(3|2) = 1.0
- p(end|2) = 1e-9
- p(1|3) = 1.0
- p(end|3) = 1e-9

The probabilities don't add up 1 because of the 1e-9 transitions to end. That doesn't
really matter. Each state just needed some transition to the end probability with a very
small probability to ensure it's possible to reach the end state from there and that it
isn't selected by beam search without a constraint.

Below is the beam search tracing for beam size 2. Any sequence below the
line is not selected by beam search. The number that comes before the sequence
is the probability of the sequence.

Step 1
1.0: [1]

Step 2
0.6: [1, 3]
0.4: [1, 2]
-----
1e-9: [1, 2, end]

Step 3
0.6: [1, 3, 1]
0.4: [1, 2, 3]
-----
0.6 * 1e-9: [1, 3, end]
0.4 * 1e-9: [1, 2, end]

Step 4
0.4:  [1, 2, 3, 1]
0.36: [1, 3, 1, 3]
-----
0.24:       [1, 3, 1, 2]
0.6 * 1e-9: [1, 3, 1, end]
0.4 * 1e-9: [1, 2, 3, end]

Step 5
0.36: [1, 3, 1, 3, 1]
0.24: [1, 2, 3, 1, 3]
-----
0.16:        [1, 2, 3, 1, 2]
0.4 * 1e-9:  [1, 2, 3, 1, end]
0.36 * 1e-9: [1, 3, 1, 3, end]
tests/nn/beam_search_test.py