Method test_forward

axlearn/audio/decoder_asr_test.py:1556–1617

(self)

        return cfg

    def test_forward(self):
        encoder_dim, decoder_dim, num_heads, vocab_size = 8, 6, 3, 20
        cfg: LASDecoderModel.Config = self._set_up_las(
            encoder_dim=encoder_dim,
            decoder_dim=decoder_dim,
            num_heads=num_heads,
            vocab_size=vocab_size,
        )
        bos_id = eos_id = cfg.decoder.eos_token_id
        pad_id = cfg.decoder.pad_token_id
        # Initialize layer parameters.
        layer: LASDecoderModel = cfg.set(name="test").instantiate(parent=None)
        prng_key = jax.random.PRNGKey(123)
        prng_key, init_key, input_key = jax.random.split(prng_key, num=3)
        layer_params = layer.initialize_parameters_recursively(init_key)

        batch_size, max_src_len = 3, 10
        target_labels = jnp.array(
            [
                [14, 8, 17, 19, 17, eos_id],  # length 5.
                [17, 4, 18, eos_id, pad_id, pad_id],  # length 3.
                [eos_id, pad_id, pad_id, pad_id, pad_id, pad_id],  # length 0.
            ]
        )
        input_ids = jnp.concatenate(
            [jnp.full([batch_size, 1], bos_id), target_labels[:, :-1]], axis=1
        )

        src_len = np.array([10, 0, 7])
        target_len = np.array([6, 4, 1])
        # [batch_size, src_len, am_dim].
        inputs = jax.random.normal(input_key, [batch_size, max_src_len, encoder_dim]) * 1000
        paddings = jnp.arange(max_src_len)[None, :] >= src_len[:, None]

        @jax.jit
        def jit_forward(input_batch):
            (loss, aux_outputs), _ = F(
                layer,
                inputs=dict(input_batch=input_batch),
                is_training=True,
                prng_key=prng_key,
                state=layer_params,
            )
            return loss, aux_outputs

        # Compute test loss.
        loss, aux_outputs = jit_forward(
            dict(
                inputs=inputs,
                paddings=paddings,
                target_labels=target_labels,
                target=dict(input_ids=input_ids),
            )
        )
        # Empty source example has weight 0.
        expected_weight = target_len * (src_len > 0)
        self.assertNestedAllClose(aux_outputs["per_example_weight"], expected_weight)
        assert_allclose(