fix(breeze): keep bf16 prompt embeddings in bf16

mlx_audio/tts/models/breeze_tts/breeze_tts.py

@@ -201,7 +201,7 @@ def __init__(self, vocab_size: int, hidden_size: int, eoi_token_index: int):
         self.eoi_token_index = eoi_token_index
 
     def __call__(self, input_ids: mx.array) -> mx.array:
-        embeds = self.weight[input_ids] * mx.array(self.weight.shape[-1] ** 0.5)
+        embeds = self.weight[input_ids] * (self.weight.shape[-1] ** 0.5)
         return mx.where(
             (input_ids == self.eoi_token_index)[..., None], self.eoi_embedding, embeds
         )

mlx_audio/tts/tests/test_breeze_tts.py

@@ -267,6 +267,22 @@ def text_ids(value):
     assert prompts == ["[S2]reference text", "[S2]<ins_bos>warm<ins_eos>target"]
 
 
+def test_bf16_prompt_embeddings_keep_the_weight_dtype(monkeypatch):
+    # A float32 scale in the text embedding used to promote every prompt
+    # activation, and with it the backbone KV cache, to float32, which made
+    # bf16 checkpoints ~3.4x slower than necessary.
+    model = Model(tiny_config())
+    model.set_dtype(mx.bfloat16)
+    monkeypatch.setattr(
+        model, "_text_ids", lambda _text: mx.array([1, 2, 3], dtype=mx.int32)
+    )
+    embeds = model._prompt_embeddings(
+        "target", voice=None, instruct=None, ref_audio=None, ref_text=None
+    )
+    assert model.text_encoder.embed_tokens.weight.dtype == mx.bfloat16
+    assert embeds.dtype == mx.bfloat16
+
+
 def test_stream_flushes_at_exact_interval_and_resets_state(monkeypatch):
     model = Model(tiny_config())