Method test_forward

axlearn/audio/decoder_asr_test.py:1155–1232

Tests that loss computation excludes empty sequence, and respects paddings.

(self, target_labels, tile_input)

    )
    @set_threefry_partitionable(False)  # TODO(Luzy): update for threefry_partitionable True
    def test_forward(self, target_labels, tile_input):
        """Tests that loss computation excludes empty sequence, and respects paddings."""
        am_dim, bos_id = 4, 1
        layer, layer_params, prng_key = self._set_up_transducer(vocab_size=20)

        # Generate inputs.
        if tile_input:
            batch_size, src_len, max_src_len = 4, 5, 10
            # [batch_size, src_len, am_dim].
            inputs = jnp.tile(
                jax.random.normal(jax.random.PRNGKey(707), [1, max_src_len, am_dim]) * 1000,
                [batch_size, 1, 1],
            )
            paddings = jnp.tile(jnp.arange(max_src_len)[None, :] >= src_len, [batch_size, 1])
            # Generate different padding data.
            pad_inputs_data = (
                jax.random.normal(jax.random.PRNGKey(124), [batch_size, max_src_len, am_dim]) * 200
            )
            # Generate inputs with the same data at non-pad positions.
            inputs = jnp.where(paddings[:, :, None], pad_inputs_data, inputs)
            assert_allclose(
                jnp.diff(inputs[:, :src_len], axis=0), jnp.zeros([batch_size - 1, src_len, am_dim])
            )
            self.assertGreater(
                jnp.abs(jnp.diff(inputs[:, src_len:], axis=0)).sum(),
                1e3,
            )
            target_labels = jnp.tile(target_labels, [batch_size, 1])
        else:
            batch_size = target_labels.shape[0]
            max_src_len = 10
            src_len = np.array([10, 0, 7])
            # [batch_size, src_len, am_dim].
            inputs = (
                jax.random.normal(jax.random.PRNGKey(311), [batch_size, max_src_len, am_dim]) * 1000
            )
            paddings = jnp.arange(max_src_len)[None, :] >= src_len[:, None]

        input_ids = jnp.concatenate(
            [jnp.full([batch_size, 1], bos_id), target_labels[:, :-1]], axis=1
        )

        @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,