def logits(self, tokens: mx.array, audio_features: mx.array) -> mx.array:
        """Perform a forward pass on the decoder and return per-token logits."""
        if tokens.shape[-1] > self.initial_token_length:
            # only need to use the last token except in the first forward pass
            tokens = tokens[:, -1:]

        logits, self.kv_cache, _ = self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache)
        return logits.astype(mx.float32)
Perform a forward pass on the decoder and return per-token logits.
whisperplus/pipelines/lightning_whisper_mlx/decoding.py