SelfAttention_MLX working

Debugged @bdeanhardt tests; they seem to be working now. <img width="1453" alt="Screenshot 2025-02-25 at 19 30 19" src="https://github.com/user-attachments/assets/9aa7a6e2-f8d9-4a1f-b67a-ea8357dbf777" />