Method test_prefix_merger

axlearn/audio/decoder_asr_test.py:738–877

(self)

    @set_threefry_partitionable(False)  # TODO(markblee): update for threefry_partitionable True
    def test_prefix_merger(self):
        # Use a small vocab_size to encourage similar prefixes.
        input_dim, vocab_size, num_decodes = 6, 3, 4
        cfg: CTCDecoderModel.Config = CTCDecoderModel.default_config().set(
            input_dim=input_dim,
            vocab_size=vocab_size,
        )
        # Initialize layer parameters.
        layer: CTCDecoderModel = cfg.set(name="test").instantiate(parent=None)
        prng_key = jax.random.PRNGKey(123)
        decode_key, init_key, input_key = jax.random.split(prng_key, num=3)
        layer_params = layer.initialize_parameters_recursively(init_key)

        batch_size, max_seq_len = 4, 8
        seq_len = jnp.array([8, 5, 4, 6])
        # [batch_size, max_seq_len, dim].
        inputs = jax.random.normal(input_key, [batch_size, max_seq_len, input_dim]) * 1000
        # [batch_size, max_seq_len].
        paddings = jnp.arange(max_seq_len) >= seq_len[:, None]

        @functools.partial(jax.jit, static_argnames=("method", "prefix_merger", "num_decodes"))
        def jit_method(inputs, prng_key, method, prefix_merger=None, num_decodes=None):
            if prefix_merger is not None:
                inputs["prefix_merger"] = prefix_merger
            if num_decodes is not None:
                inputs["num_decodes"] = num_decodes
            outputs, _ = F(
                layer,
                inputs=inputs,
                is_training=True,
                prng_key=prng_key,
                state=layer_params,
                method=method,
            )
            return outputs

        input_batch = dict(input_batch=dict(inputs=inputs, paddings=paddings))

        # Decode without merging.
        beam_search_outputs: DecodeOutputs = jit_method(
            input_batch,
            prng_key=decode_key,
            method="beam_search_decode",
            num_decodes=num_decodes,
        )
        self.assertNestedEqual(
            jnp.array(
                [
                    [
                        # Without merging prefixes, sequences 0 and 1 are duplicates.
                        [1, 2, 1, 2, 1, 0, 0, 0],
                        [1, 2, 1, 2, 1, 0, 0, 0],
                        [1, 2, 1, 2, 1, 2, 0, 0],
                        [1, 2, 2, 1, 0, 0, 0, 0],
                    ],
                    [
                        # Sequence 0 and 3 are duplicates.
                        [2, 1, 1, 2, 0, 0, 0, 0],