ci(mlx): lazy mx.transpose repro + keep long build alive

Probe gather_mm(sorted=True) on an unevaluated transpose op (mirrors
ET TransposeNode). Disable cancel-in-progress / ignore repro-only
paths so the fast job can run beside the Vid-vs-literal build.

Author: AxelNoun

.github/workflows/mlx-fork-check.yml

@@ -25,10 +25,14 @@ on:
   push:
     branches:
       - ci-mlx-check
+    paths-ignore:
+      - mlx_gather_mm_repro.py
+      - .github/workflows/mlx-only-repro.yml
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  # false: a repro-only push must not cancel an in-flight ExecuTorch build
+  cancel-in-progress: false
 
 permissions:
   contents: read

mlx_gather_mm_repro.py

@@ -17,29 +17,37 @@
 E, K, NOUT = 4, 64, 128
 
 
-def run_case(L, dtype, idx_np, a_ndim, label, b_strided=False):
+def run_case(L, dtype, idx_np, a_ndim, label, b_mode="contig"):
     """Compare référence / sorted=False / sorted=True et localise les lignes fausses.
 
-    b_strided=True reproduit EXACTEMENT ce que fait GatherMmModel dans
-    test_ops.py : le buffer est stocké en [E, out, in] et transposé à la volée
-    en [E, in, out] -> vue NON CONTIGUE (strides échangés sur les 2 derniers
-    axes). Le chemin de production (SwitchLinear.pack) fait au contraire
-    .transpose(-1,-2).contiguous(), donc b y est toujours contigu.
+    b_mode:
+      contig          — b = mx.array([E,K,N]) contigu (baseline)
+      swapaxes        — vue non contiguë via mx.array(...).swapaxes (déjà OK)
+      lazy_transpose  — mx.transpose(mx.array(b_src), (0,2,1)) SANS eval avant
+                        gather_mm : imite le TransposeNode du graphe ET (array
+                        paresseux issu d'une op, pas une vue construite à la main)
     """
     rng = np.random.default_rng(0)
     a_np = rng.standard_normal((L, 1, K)).astype(np.float32)
     b_np = rng.standard_normal((E, K, NOUT)).astype(np.float32)
 
     a = mx.array(a_np if a_ndim == 3 else a_np.reshape(L, K)).astype(dtype)
-    if b_strided:
-        # mêmes valeurs logiques, stockage transposé
+    if b_mode == "contig":
+        b = mx.array(b_np).astype(dtype)
+    elif b_mode == "swapaxes":
+        # mêmes valeurs logiques, stockage transposé → vue strides échangés
+        b_src = np.ascontiguousarray(b_np.transpose(0, 2, 1))  # [E, NOUT, K]
+        b = mx.array(b_src).astype(dtype).swapaxes(-1, -2)  # [E, K, NOUT]
+    elif b_mode == "lazy_transpose":
+        # Comme le TransposeNode ET : op MLX, pas de mx.eval avant gather_mm.
         b_src = np.ascontiguousarray(b_np.transpose(0, 2, 1))  # [E, NOUT, K]
-        b = mx.array(b_src).astype(dtype).swapaxes(-1, -2)  # [E, K, NOUT] non contigu
+        b = mx.transpose(mx.array(b_src).astype(dtype), (0, 2, 1))  # [E, K, NOUT]
     else:
-        b = mx.array(b_np).astype(dtype)
+        raise ValueError(f"unknown b_mode={b_mode!r}")
     idx = mx.array(idx_np.astype(np.uint32))
 
-    # Référence explicite : take puis matmul (ce que fait l'eager PyTorch).
+    # Référence + gather : un seul mx.eval à la fin, pour ne pas matérialiser
+    # le lazy_transpose avant le kernel.
     ref = mx.matmul(a, mx.take(b, idx, axis=0))
     got_unsorted = mx.gather_mm(a, b, None, idx, sorted_indices=False)
     got_sorted = mx.gather_mm(a, b, None, idx, sorted_indices=True)
@@ -48,7 +56,7 @@ def run_case(L, dtype, idx_np, a_ndim, label, b_strided=False):
     d_uns = float(mx.max(mx.abs(got_unsorted.astype(mx.float32) - ref.astype(mx.float32))))
     d_srt = float(mx.max(mx.abs(got_sorted.astype(mx.float32) - ref.astype(mx.float32))))
 
-    print(f"\n--- {label} | L={L} a.ndim={a_ndim} dtype={dtype} idx={idx_np.tolist()}")
+    print(f"\n--- {label} | L={L} a.ndim={a_ndim} dtype={dtype} b_mode={b_mode} idx={idx_np.tolist()}")
     print(f"    sorted=False  max_diff = {d_uns:.6e}   {'OK' if d_uns < 1e-2 else 'DIVERGE'}")
     print(f"    sorted=True   max_diff = {d_srt:.6e}   {'OK' if d_srt < 1e-2 else 'DIVERGE'}")
 
@@ -78,10 +86,10 @@ def run_case(L, dtype, idx_np, a_ndim, label, b_strided=False):
 except Exception:
     pass
 print(
-    "ATTENTION : le backend ExecuTorch NE lie PAS ce paquet PyPI.\n"
-    "Il compile le sous-module epingle backends/mlx/third-party/mlx\n"
-    "(commit 7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247 sur main).\n"
-    "Un OK ici ne dit rien de la version reellement liee dans le build.\n"
+    "ATTENTION : le backend ExecuTorch lie le sous-module\n"
+    "backends/mlx/third-party/mlx @ 7a1d4f5c… (= tag v0.32.0, meme SHA\n"
+    "que pip install mlx==0.32.0). Un OK ici EST comparable au runtime ET\n"
+    "pour la version MLX ; l'ecart restant est la construction des arrays.\n"
 )
 ok = []
 
@@ -101,12 +109,24 @@ def run_case(L, dtype, idx_np, a_ndim, label, b_strided=False):
 #    layout [L, 1, K] (celui du chemin MoE) qui n'est pas géré.
 ok.append(run_case(8, mx.float32, np.array([0, 0, 0, 1, 1, 2, 3, 3]), 2, "4 segments, a 2-D"))
 
-# 5. LE CAS DU MODELE DE TEST : b non contigu (weight.transpose sans .contiguous()).
-#    C'est le seul écart entre GatherMmModel et le chemin de production.
+# 5. Vue non contiguë via swapaxes (deja innocenté).
+ok.append(run_case(2, mx.float32, np.array([1, 3]), 3,
+                   "cas du test, b swapaxes", b_mode="swapaxes"))
+ok.append(run_case(8, mx.float32, np.array([0, 0, 0, 1, 1, 2, 3, 3]), 3,
+                   "4 segments, b swapaxes", b_mode="swapaxes"))
+
+# 5b. LE CAS QUI MANQUAIT : transpose LAZY (op MLX, pas de eval avant gather_mm).
+#     C'est ce que produit TransposeNode dans le graphe ET — distinct d'une
+#     vue swapaxes construite a la main. Si ca diverge : bug MLX hors PR.
 ok.append(run_case(2, mx.float32, np.array([1, 3]), 3,
-                   "cas du test, b NON CONTIGU", b_strided=True))
+                   "cas du test, b LAZY transpose", b_mode="lazy_transpose"))
+ok.append(run_case(2, mx.bfloat16, np.array([1, 3]), 3,
+                   "cas du test bf16, b LAZY transpose", b_mode="lazy_transpose"))
 ok.append(run_case(8, mx.float32, np.array([0, 0, 0, 1, 1, 2, 3, 3]), 3,
-                   "4 segments, b NON CONTIGU", b_strided=True))
+                   "4 segments, b LAZY transpose", b_mode="lazy_transpose"))
+# Contre-epreuve : lazy transpose + sorted=False doit rester OK.
+ok.append(run_case(2, mx.float32, np.array([1, 3]), 3,
+                   "lazy transpose (unsorted indices contract)", b_mode="lazy_transpose"))
 
 # 7. BALAYAGE DU DTYPE DES INDICES.
 #    _gather_mm_handler passe rhs_indices tel quel, SANS cast (verifie dans
@@ -137,6 +157,30 @@ def run_case(L, dtype, idx_np, a_ndim, label, b_strided=False):
     except Exception as e:  # dtype refuse par MLX
         print(f"    idx {np.dtype(_idt).name:8s} | rejete : {type(e).__name__}: {e}")
 
+# 7b. Meme balayage dtype sur b LAZY transpose (le combo du graphe ET).
+print("\n--- balayage dtype + b LAZY transpose (sorted=True, idx=[1,3])")
+_rng2 = np.random.default_rng(0)
+_a2 = mx.array(_rng2.standard_normal((2, 1, K)).astype(np.float32))
+_b2_np = _rng2.standard_normal((E, K, NOUT)).astype(np.float32)
+_b2_src = np.ascontiguousarray(_b2_np.transpose(0, 2, 1))
+_b2 = mx.transpose(mx.array(_b2_src), (0, 2, 1))  # lazy, no eval
+_idx_ref2 = mx.array(np.array([1, 3], dtype=np.uint32))
+_ref2 = mx.matmul(_a2, mx.take(_b2, _idx_ref2, axis=0))
+for _idt in (np.uint32, np.int32, np.int64):
+    try:
+        _idx = mx.array(np.array([1, 3], dtype=_idt))
+        _got = mx.gather_mm(_a2, _b2, None, _idx, sorted_indices=True)
+        mx.eval(_ref2, _got)  # first eval materializes both
+        _d = float(mx.max(mx.abs(_got - _ref2)))
+        flag = "OK" if _d < 1e-2 else "DIVERGE <<<<"
+        print(f"    idx {np.dtype(_idt).name:8s} + lazy_transpose | sorted=True {_d:.3e} {flag}")
+        ok.append(_d < 1e-2)
+        # rebuild lazy b for next iter (previous eval materialized it)
+        _b2 = mx.transpose(mx.array(_b2_src), (0, 2, 1))
+        _ref2 = mx.matmul(_a2, mx.take(_b2, _idx_ref2, axis=0))
+    except Exception as e:
+        print(f"    idx {np.dtype(_idt).name:8s} + lazy_transpose | rejete : {type(e).__name__}: {e}")
+
 # 6. Contre-épreuve quantifiée : gather_qmm sur le même layout (il passe en CI).
 print("\n--- controle gather_qmm (meme layout)")
 rng = np.random.default_rng(0)