Method test_beam_search_with_path_merger

axlearn/common/decoding_test.py:727–812

(self)

        )

    def test_beam_search_with_path_merger(self):
        batch_size, num_decodes = 1, 4
        max_decode_len = 8
        vocab_size = 4
        num_expected_tokens = 5
        cache = dict(step=jnp.zeros([]))

        def tokens_to_scores(tokens, cache):
            del tokens
            step = cache["step"]
            has_enough_tokens = step >= num_expected_tokens
            eps = 1e-2 + 1e-3 * step
            continue_log_probs = jax.nn.log_softmax(
                # A mostly uniform distribution among tokens except for PAD(0) and EOS(1) with
                # slight preference towards lower ids.
                #
                # [NEG_INF, NEG_INF, 0, -eps, -2 * eps, ...].
                jnp.pad(
                    jnp.arange(0, vocab_size - 2, dtype=jnp.float32) * -eps,
                    ((2, 0)),
                    constant_values=NEG_INF,
                )
            )
            # Force EOS=1: [NEG_INF, 0, NEG_INF, NEG_INF, ...].
            finish_log_probs = (1 - jax.nn.one_hot(1, vocab_size)) * NEG_INF
            log_probs = (
                has_enough_tokens * finish_log_probs + (1 - has_enough_tokens) * continue_log_probs
            )
            return log_probs, dict(step=cache["step"] + 1)

        inputs = jnp.zeros([batch_size, max_decode_len], dtype=jnp.int32)
        kwargs = dict(
            inputs=inputs,
            time_step=decoding.infer_initial_time_step(inputs, pad_id=0),
            cache=cache,
            tokens_to_scores=tokens_to_scores,
            eos_id=1,
            num_decodes=num_decodes,
        )
        outputs_without_merger = decoding.beam_search_decode(**kwargs)

        prefix_merger = _TokenSumPrefixMerger()
        outputs_with_merger = decoding.beam_search_decode(**kwargs, prefix_merger=prefix_merger)

        logging.info(
            "scores_without_merger=%s scores_with_merger=%s",
            outputs_without_merger.scores,
            outputs_with_merger.scores,
        )

        np.testing.assert_array_equal(
            jnp.asarray(
                [
                    [
                        # Without prefix merging, we prefer sequences with mostly 2's.
                        [2, 2, 2, 2, 2, 1, 0, 0],
                        [3, 2, 2, 2, 2, 1, 0, 0],
                        [2, 3, 2, 2, 2, 1, 0, 0],