Method test_forward

axlearn/audio/decoder_asr_test.py:439–468

(self, blank_id)

    @parameterized.parameters([0, 1])
    def test_forward(self, blank_id):
        vocab_size = 20
        golden = load_golden("axlearn.audio.decoder_asr_test", f"test_forward_blank{blank_id}")
        # Load golden inputs.
        logits = jnp.asarray(golden["inputs"]["logits"])
        input_lengths = jnp.asarray(golden["inputs"]["input_lengths"])
        target_labels = jnp.asarray(golden["inputs"]["target_labels"])
        max_seq_len = logits.shape[1]

        # Compute paddings and target_paddings from lengths.
        paddings = (jnp.arange(max_seq_len) >= input_lengths[:, None]).astype(jnp.float32)
        target_paddings = jnp.logical_or(vocab_size <= target_labels, target_labels < 0)

        # Compute CTC loss using optax (same as CTCDecoderModel.forward).
        per_example_loss = optax.ctc_loss(
            logits=logits,
            logit_paddings=paddings,
            labels=target_labels,
            label_paddings=target_paddings,
            blank_id=blank_id,
        )

        # Mask invalid sequences (optax returns large values; torch zero_infinity returns 0).
        per_example_weight = _is_valid_ctc_seq(
            paddings=paddings, target_labels=target_labels, target_paddings=target_paddings
        ).astype(jnp.float32)

        # Compare against golden reference (torch CTC loss, weighted by validity).
        ref_per_example_loss = jnp.asarray(golden["outputs"]["per_example_loss"])
        assert_allclose(ref_per_example_loss, per_example_loss * per_example_weight)

    def _check_summary(
        self, summary_collection: dict[str, Any], name: str, value: Union[Tensor, WeightedSummary]