2b1a3c1312
Byte-parity with pyannote-PyTorch reference (cosine 0.763718 identical at 6 decimals on 200 cross-window slot pairs). 2.5x faster than pyannote-MPS on Apple Silicon native. Extracted from gitea.tavportal.com/olivier/MLX_CONVERTOR commit 5f9eafa.
12 lines
415 B
Python
12 lines
415 B
Python
import mlx.core as mx
|
|
from pyannote_diarization_3_1_mlx._bilstm import BiLSTM4
|
|
|
|
|
|
def test_bilstm_output_shape():
|
|
# input (B, T, hidden_in) — pyannote feeds 60-channel sincnet output
|
|
# transposed to (B, T, 60). hidden=128, bidirectional → 256 out.
|
|
net = BiLSTM4(input_size=60, hidden_size=128)
|
|
x = mx.zeros((1, 589, 60))
|
|
out = net(x)
|
|
assert out.shape == (1, 589, 256), f"got {out.shape}"
|