def sanitize_lstm_weights(key: str, state_dict: mx.array) -> dict:
    """Convert PyTorch LSTM weight keys to MLX LSTM weight keys."""
    base_key = key.rsplit(".", 1)[0]

    # Mapping of PyTorch LSTM weight suffixes to MLX weight names
    weight_map = {
        "weight_ih_l0_reverse": "Wx_backward",
        "weight_hh_l0_reverse": "Wh_backward",
        "bias_ih_l0_reverse": "bias_ih_backward",
        "bias_hh_l0_reverse": "bias_hh_backward",
        "weight_ih_l0": "Wx_forward",
        "weight_hh_l0": "Wh_forward",
        "bias_ih_l0": "bias_ih_forward",
        "bias_hh_l0": "bias_hh_forward",
    }

    for suffix, new_suffix in weight_map.items():
        if key.endswith(suffix):
            return {f"{base_key}.{new_suffix}": state_dict}

    return {key: state_dict}
Convert PyTorch LSTM weight keys to MLX LSTM weight keys.
mlx_audio/tts/models/kokoro/kokoro.py