import mlx.core as mx

from ark_asr_mlx.config import AudioEncoderConfig
from ark_asr_mlx.encoder import WhisperRoPEEncoder, apply_interleaved_rope


def test_interleaved_rope_preserves_shape() -> None:
    values = mx.ones((1, 2, 5, 8))
    output = apply_interleaved_rope(values, rotary_dim=4)
    assert output.shape == values.shape


def test_audio_encoder_downsamples_time_by_two() -> None:
    config = AudioEncoderConfig(
        d_model=8,
        encoder_attention_heads=2,
        encoder_ffn_dim=16,
        encoder_layers=1,
        num_mel_bins=4,
    )
    encoder = WhisperRoPEEncoder(config)
    output = encoder(mx.zeros((1, 4, 10)))
    mx.eval(output)
    assert output.shape == (1, 5, 8)