fix: guard conv-weight PyTorch->MLX transpose with layout detection

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
+## [Unreleased]
+
+### Fixed
+- **Double-transpose crash on mlx-community conv checkpoints (Gemma 4 audio, phi4mm patch-embed, nemotron audio, RT-DETRv2).** Several weight-sanitizer functions transposed conv weights from PyTorch `[out, in, kH, kW]` to MLX channel-last `[out, kH, kW, in]` unconditionally. Pre-converted mlx-community checkpoints already store these weights in channel-last order, so the unconditional transpose double-converted them and produced a corrupted shape. The confirmed crash: loading `mlx-community/gemma-4-e4b-it-qat-4bit` turned the audio subsample conv weight `[128, 3, 3, 1]` into `[128, 3, 1, 3]`, which MLX conv2d rejected because the input C_in=1 did not match the weight C_in=3. All four affected sanitizers now check the tensor shape before transposing: `conv2d_weight_is_channel_last` (already-MLX `[out, kH, kW, in]` skips; PyTorch `[out, in, kH, kW]` transposes) and `conv1d_weight_is_channel_last` (depthwise-only; MLX `[out, kW, 1]` skips; PyTorch `[out, 1, kW]` transposes). Both predicates are idempotent. Resolves #428.
+
 ## [v0.3.3] - 2026-06-23
 
 ### Added

src/loading/mod.rs

@@ -88,6 +88,53 @@ pub(super) fn parse_model_config<T: DeserializeOwned>(config_str: &str) -> Resul
     serde_json::from_str(config_str).map_err(|e| anyhow::anyhow!("Failed to parse config: {}", e))
 }
 
+/// Layout-detection guard for conv weights loaded from a checkpoint.
+///
+/// Several loaders transpose conv weights from PyTorch to MLX channel-last
+/// layout while sanitizing a checkpoint. Pre-converted `mlx-community`
+/// checkpoints already store these weights channel-last, so an unconditional
+/// transpose double-converts and corrupts the shape (see issue #428). These
+/// helpers detect the layout from the shape so the transpose fires only for
+/// genuine PyTorch-layout weights, which is safe in both directions: a true
+/// PyTorch weight is still transposed and an already-MLX weight is left as-is.
+///
+/// Returns `true` when a 4D conv2d weight is already in MLX channel-last
+/// `[out, kH, kW, in]` layout (skip the transpose). PyTorch layout is
+/// `[out, in, kH, kW]`.
+///
+/// Heuristic: in channel-last layout the output-channel count is the leading
+/// dim and dominates both spatial dims, and the two kernel dims are equal
+/// (square kernels are the norm for these patch/subsample/backbone convs).
+/// PyTorch layout puts `in` in dim 1, which generally breaks `dim1 == dim2`
+/// (e.g. `[128, 1, 3, 3]` has `dim1=1 != dim2=3`) or `out >= dim1` once the
+/// in-channel count exceeds the output count. This mirrors the validated
+/// `should_transpose_phi3_patch_embedding` predicate.
+#[must_use]
+pub(crate) fn conv2d_weight_is_channel_last(shape: &[i32]) -> bool {
+    shape.len() == 4 && shape[0] >= shape[1] && shape[0] >= shape[2] && shape[1] == shape[2]
+}
+
+/// Returns `true` when a 3D depthwise conv1d weight is already in MLX
+/// channel-last `[out, kW, in]` layout (skip the transpose). PyTorch layout is
+/// `[out, in, kW]`.
+///
+/// This predicate is valid only for depthwise kernels, where the in-channel
+/// count per group is `1`. In that case channel-last `[out, kW, 1]` carries the
+/// `1` in the trailing dim, while PyTorch `[out, 1, kW]` carries it in the
+/// middle dim. Detecting `shape[2] == 1 && shape[1] > 1` therefore distinguishes
+/// the two depthwise layouts unambiguously: the confirmed Gemma 4 audio kernel
+/// `[1024, 5, 1]` is recognized as already-MLX (skip), and its PyTorch form
+/// `[1024, 1, 5]` as needing the transpose.
+///
+/// It does NOT generalize to pointwise conv1d (`kernel == 1`): PyTorch pointwise
+/// `[out, in, 1]` and MLX depthwise `[out, kW, 1]` share the same shape
+/// signature and cannot be told apart from shape alone. Callers that may receive
+/// pointwise weights must scope this guard to depthwise keys.
+#[must_use]
+pub(crate) fn conv1d_weight_is_channel_last(shape: &[i32]) -> bool {
+    shape.len() == 3 && shape[2] == 1 && shape[1] > 1
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum Qwen35VlmKind {
     Dense,

src/loading/tests.rs

@@ -13,8 +13,9 @@
 // limitations under the License.
 
 use super::{
-    Qwen35VlmKind, model_path_str, parse_eos_token_ids, qwen35_vlm_kind, read_eos_token_ids,
-    require_qwen35_vlm_kind, resolve_model_dir,
+    Qwen35VlmKind, conv1d_weight_is_channel_last, conv2d_weight_is_channel_last, model_path_str,
+    parse_eos_token_ids, qwen35_vlm_kind, read_eos_token_ids, require_qwen35_vlm_kind,
+    resolve_model_dir,
 };
 use crate::model_metadata::{
     DirectoryLoadRoute, ModelCapabilities, ModelKind, ModelLoadPolicy, WeightLoadRoute,
@@ -347,3 +348,29 @@ fn model_load_policy_handles_mistral3_mistral4_wrapper_config() {
     );
     assert_eq!(policy.weight_route, Some(WeightLoadRoute::LlamaFamily));
 }
+
+#[test]
+fn conv2d_weight_is_channel_last_detects_mlx_layout() {
+    // MLX channel-last [out, kH, kW, in]: out dominates and kernels are square.
+    assert!(conv2d_weight_is_channel_last(&[128, 3, 3, 1]));
+    assert!(conv2d_weight_is_channel_last(&[32, 3, 3, 128]));
+    // PyTorch [out, in, kH, kW]: in breaks dim1 == dim2 (or out >= dim1).
+    assert!(!conv2d_weight_is_channel_last(&[128, 1, 3, 3]));
+    assert!(!conv2d_weight_is_channel_last(&[32, 128, 3, 3]));
+    // Non-4D shapes are never treated as conv2d channel-last.
+    assert!(!conv2d_weight_is_channel_last(&[128, 3, 3]));
+    assert!(!conv2d_weight_is_channel_last(&[4, 4]));
+}
+
+#[test]
+fn conv1d_weight_is_channel_last_detects_depthwise_mlx_layout() {
+    // MLX depthwise [out, kW, in=1]: trailing in dim is 1.
+    assert!(conv1d_weight_is_channel_last(&[1024, 5, 1]));
+    // PyTorch depthwise [out, in=1, kW]: middle in dim is 1, kW > 1.
+    assert!(!conv1d_weight_is_channel_last(&[1024, 1, 5]));
+    // A regular conv1d [out, in, kW] with in, kW > 1 is treated as PyTorch.
+    assert!(!conv1d_weight_is_channel_last(&[4, 8, 3]));
+    // Non-3D shapes are never treated as depthwise conv1d channel-last.
+    assert!(!conv1d_weight_is_channel_last(&[1024, 5, 1, 1]));
+    assert!(!conv1d_weight_is_channel_last(&[4, 4]));
+}

src/loading/vlm_gemma.rs

@@ -280,17 +280,25 @@ fn sanitize_gemma4_audio_weights(weights: &mut mlxcel_core::weights::WeightMap)
             mlxcel_core::array_shape(w)
         };
 
-        // Conv2d: PyTorch [out, in, kH, kW] -> MLX [out, kH, kW, in]
+        // Conv2d: PyTorch [out, in, kH, kW] -> MLX [out, kH, kW, in].
+        // Skip already-channel-last checkpoints to avoid double-converting the
+        // shape (issue #428): the mlx-community gemma4 weights ship channel-last.
         if key.contains("subsample_conv_projection")
             && key.contains("conv.weight")
             && shape.len() == 4
+            && !crate::loading::conv2d_weight_is_channel_last(&shape)
         {
             let w = weights.remove(&key).unwrap();
             let transposed = mlxcel_core::transpose_axes(&w, &[0, 2, 3, 1]);
             weights.insert(key, transposed);
         }
-        // Conv1d depthwise: PyTorch [out, in, kW] -> MLX [out, kW, in]
-        else if key.contains("depthwise_conv1d.weight") && shape.len() == 3 {
+        // Conv1d depthwise: PyTorch [out, in, kW] -> MLX [out, kW, in].
+        // The Conformer lconv1d is depthwise (in == 1), so the channel-last
+        // form is [out, kW, 1]; skip it when already converted (issue #428).
+        else if key.contains("depthwise_conv1d.weight")
+            && shape.len() == 3
+            && !crate::loading::conv1d_weight_is_channel_last(&shape)
+        {
             let w = weights.remove(&key).unwrap();
             let transposed = mlxcel_core::transpose_axes(&w, &[0, 2, 1]);
             weights.insert(key, transposed);

src/loading/vlm_gemma_tests.rs

@@ -14,7 +14,7 @@
 
 use super::{
     gemma3n_language_model_prefix, gemma3n_metadata, gemma3n_needs_conv_transpose,
-    sanitize_gemma3n_weights,
+    sanitize_gemma3n_weights, sanitize_gemma4_audio_weights,
 };
 use mlxcel_core::dtype;
 use mlxcel_core::weights::WeightMap;
@@ -109,3 +109,101 @@ fn sanitize_gemma3n_weights_strips_model_prefix_and_transposes_conv_weights() {
         vec![8, 3, 3, 16]
     );
 }
+
+fn audio_conv_shape(weights: &WeightMap, key: &str) -> Vec<i32> {
+    mlxcel_core::array_shape(weights.get(key).unwrap())
+}
+
+const LAYER0_CONV: &str = "audio_tower.subsample_conv_projection.layer0.conv.weight";
+const LAYER1_CONV: &str = "audio_tower.subsample_conv_projection.layer1.conv.weight";
+const LCONV1D: &str = "audio_tower.layers.0.lconv1d.depthwise_conv1d.weight";
+
+#[test]
+fn sanitize_gemma4_audio_weights_preserves_channel_last_checkpoint() {
+    // The mlx-community/gemma-4-e4b-it-qat-4bit checkpoint stores these conv
+    // weights already in MLX channel-last layout. The sanitizer must leave them
+    // untouched so the audio subsample conv receives a C_in=1 weight matching
+    // the C_in=1 input (regression for issue #428).
+    let mut weights = WeightMap::new();
+    weights.insert(
+        LAYER0_CONV.to_string(),
+        mlxcel_core::ones(&[128, 3, 3, 1], dtype::FLOAT32),
+    );
+    weights.insert(
+        LAYER1_CONV.to_string(),
+        mlxcel_core::ones(&[32, 3, 3, 128], dtype::FLOAT32),
+    );
+    weights.insert(
+        LCONV1D.to_string(),
+        mlxcel_core::ones(&[1024, 5, 1], dtype::FLOAT32),
+    );
+
+    sanitize_gemma4_audio_weights(&mut weights);
+
+    assert_eq!(audio_conv_shape(&weights, LAYER0_CONV), vec![128, 3, 3, 1]);
+    assert_eq!(audio_conv_shape(&weights, LAYER1_CONV), vec![32, 3, 3, 128]);
+    assert_eq!(audio_conv_shape(&weights, LCONV1D), vec![1024, 5, 1]);
+}
+
+#[test]
+fn sanitize_gemma4_audio_weights_transposes_pytorch_layout() {
+    // Synthetic PyTorch-layout weights: conv2d [out, in, kH, kW] and depthwise
+    // conv1d [out, in=1, kW]. Both must be transposed to channel-last.
+    let mut weights = WeightMap::new();
+    weights.insert(
+        LAYER0_CONV.to_string(),
+        mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32),
+    );
+    weights.insert(
+        LCONV1D.to_string(),
+        mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32),
+    );
+
+    sanitize_gemma4_audio_weights(&mut weights);
+
+    // conv2d [128, 1, 3, 3] --transpose[0,2,3,1]--> [128, 3, 3, 1].
+    assert_eq!(audio_conv_shape(&weights, LAYER0_CONV), vec![128, 3, 3, 1]);
+    // depthwise conv1d [1024, 1, 5] --transpose[0,2,1]--> [1024, 5, 1].
+    assert_eq!(audio_conv_shape(&weights, LCONV1D), vec![1024, 5, 1]);
+}
+
+#[test]
+fn sanitize_gemma4_audio_weights_is_idempotent() {
+    // Running the sanitizer twice must equal running it once: the first pass
+    // converts the PyTorch-layout weights to channel-last, and the second pass
+    // must detect that layout and leave them unchanged.
+    let mut once = WeightMap::new();
+    once.insert(
+        LAYER0_CONV.to_string(),
+        mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32),
+    );
+    once.insert(
+        LCONV1D.to_string(),
+        mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32),
+    );
+    sanitize_gemma4_audio_weights(&mut once);
+
+    let mut twice = WeightMap::new();
+    twice.insert(
+        LAYER0_CONV.to_string(),
+        mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32),
+    );
+    twice.insert(
+        LCONV1D.to_string(),
+        mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32),
+    );
+    sanitize_gemma4_audio_weights(&mut twice);
+    sanitize_gemma4_audio_weights(&mut twice);
+
+    assert_eq!(
+        audio_conv_shape(&once, LAYER0_CONV),
+        audio_conv_shape(&twice, LAYER0_CONV)
+    );
+    assert_eq!(
+        audio_conv_shape(&once, LCONV1D),
+        audio_conv_shape(&twice, LCONV1D)
+    );
+    // And the idempotent result is the channel-last layout.
+    assert_eq!(audio_conv_shape(&twice, LAYER0_CONV), vec![128, 3, 3, 1]);
+    assert_eq!(audio_conv_shape(&twice, LCONV1D), vec![1024, 5, 1]);
+}