File size: 709 Bytes
bd15125 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | 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)
|