From 34659fff2d1983cf69dc4e1dc9bd596aee64ed8a Mon Sep 17 00:00:00 2001
From: Scott Roy <scroy@meta.com>
Date: Mon, 21 Sep 2026 10:10:12 -0700
Subject: [PATCH 1/4] up

---
 backends/mlx/builder/op_helpers.py |  9 ++--
 backends/mlx/ops.py                | 11 +++--
 backends/mlx/test/test_ops.py      | 75 ++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 7 deletions(-)

diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py
index 48a2a0376ed..2ab88da4dbb 100644
--- a/backends/mlx/builder/op_helpers.py
+++ b/backends/mlx/builder/op_helpers.py
@@ -187,6 +187,7 @@ def emit_shape(
     slot: Slot,
     *,
     end_dim: "Optional[int]" = None,
+    dim_offset: int = 0,
 ) -> "list[IntOrVid]":
     """Return the shape of ``node`` as a list of ``IntOrVid``.
 
@@ -197,11 +198,13 @@ def emit_shape(
     Args:
         P: program builder.
         node: FX node whose shape to walk (must have ``meta['val']``).
-        slot: slot corresponding to ``node`` (used as the
-            ``SymSize`` source for any dynamic dim).
+        slot: tensor slot used as the ``SymSize`` source for dynamic dims.
         end_dim: stop index (exclusive). ``None`` means the full ndim.
             Negative values index from the end (e.g. ``-1`` is "all
             leading dims, drop the last").
+        dim_offset: offset added to dynamic dimension indices when reading
+            ``slot``. Use when its runtime axes differ from the metadata axes;
+            static dimensions still use the metadata values.
 
     Returns:
         ``list[IntOrVid]`` of length ``end_dim`` (after normalization).
@@ -228,7 +231,7 @@ def emit_shape(
             P.emit(
                 SymSizeNode(
                     a=P.slot_to_tid(slot),
-                    dim=dim_idx,
+                    dim=dim_idx + dim_offset,
                     out=P.slot_to_vid(d_val),
                 )
             )
diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index af9c4a1c821..3dff908ede5 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -2466,8 +2466,7 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     indices = [P.slot_to_tid(idx) for _, idx in non_none]
     axes = [i for i, _ in non_none]
 
-    # slice_sizes: 1 for indexed axes, full dim size for non-indexed axes
-    # Use int() to handle SymInt values from dynamic shapes
+    # slice_sizes: 1 for indexed axes, full static size for non-indexed axes.
     indexed_axes = set(axes)
     slice_sizes = []
     for dim in range(x_ndim):
@@ -2495,13 +2494,17 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     )
 
     # Reshape to match aten.index.Tensor output shape, which strips the
-    # trailing dimensions introduced by gather's slice_sizes
+    # trailing dimensions introduced by gather's slice_sizes.
     out_meta = n.meta.get("val")
     if out_meta is None:
         raise ValueError(
             "aten.index.Tensor: output shape metadata required for reshape after gather"
         )
-    out_shape = [P.to_int_or_vid(int(d)) for d in out_meta.shape]
+    # Non-indexed sizes are static above, so symbolic output sizes belong to
+    # the broadcast index shape, which leads the gather result. Read them at
+    # runtime instead of adding specialization guards through int(SymInt).
+    leading_dims = axes[0] if axes == list(range(axes[0], axes[-1] + 1)) else 0
+    out_shape = emit_shape(P, n, gather_slot, dim_offset=-leading_dims)
 
     out = P.make_or_get_slot(n)
     P.emit(
diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py
index 6595b37bd84..ad460d49cf6 100644
--- a/backends/mlx/test/test_ops.py
+++ b/backends/mlx/test/test_ops.py
@@ -4518,6 +4518,81 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (x, *indices)
 
 
+class DynamicAdvancedIndexModel(nn.Module):
+    def __init__(self, leading_singleton: bool):
+        super().__init__()
+        self.leading_singleton = leading_singleton
+
+    def forward(self, x, rows, columns):
+        if self.leading_singleton:
+            return x[:, rows, columns]
+        return x[rows, columns]
+
+
+@register_test
+class DynamicAdvancedIndexTest(OpTestCase):
+    """Broadcast index dimensions must remain dynamic through gather/reshape."""
+
+    name = "dynamic_advanced_index"
+    rtol = 1e-4
+    atol = 1e-4
+
+    def __init__(
+        self,
+        trailing_dim: bool = False,
+        leading_singleton: bool = False,
+        test_rows: int = 4,
+        test_columns: int = 5,
+    ):
+        self.trailing_dim = trailing_dim
+        self.leading_singleton = leading_singleton
+        self.test_rows = test_rows
+        self.test_columns = test_columns
+        self.name = (
+            f"dynamic_advanced_index_tail{trailing_dim}_leading{leading_singleton}"
+            f"_runtime{test_rows}x{test_columns}"
+        )
+
+    @classmethod
+    def get_test_configs(cls) -> List["DynamicAdvancedIndexTest"]:
+        return [
+            cls(
+                trailing_dim=trailing,
+                leading_singleton=leading,
+                test_rows=rows,
+                test_columns=columns,
+            )
+            for trailing, leading in ((False, False), (True, False), (True, True))
+            for rows, columns in ((4, 5), (3, 2))
+        ]
+
+    def create_model(self) -> nn.Module:
+        return DynamicAdvancedIndexModel(self.leading_singleton)
+
+    def _inputs(self, rows, columns):
+        shape = (6, 7) + ((4,) if self.trailing_dim else ())
+        if self.leading_singleton:
+            shape = (1,) + shape
+        return (
+            torch.randn(shape),
+            torch.arange(rows).reshape(-1, 1),
+            torch.arange(columns).reshape(1, -1),
+        )
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return self._inputs(2, 3)
+
+    def create_test_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return self._inputs(self.test_rows, self.test_columns)
+
+    def get_dynamic_shapes(self) -> Optional[Dict]:
+        return {
+            "x": None,
+            "rows": {0: Dim("rows", min=2, max=4)},
+            "columns": {1: Dim("columns", min=2, max=5)},
+        }
+
+
 class IndexUpdateModel(nn.Module):
     """Model that performs index_copy on a mutable buffer.
 

From e12595ef4c403dd5c3447ef6286265d0700f9a3a Mon Sep 17 00:00:00 2001
From: Scott Roy <scroy@meta.com>
Date: Mon, 21 Sep 2026 13:18:47 -0700
Subject: [PATCH 2/4] up

---
 backends/mlx/ops.py           | 26 +++++++++-
 backends/mlx/test/test_ops.py | 92 ++++++++++++++++++++++++++++++-----
 2 files changed, 106 insertions(+), 12 deletions(-)

diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 3dff908ede5..8e6e7b6ea3d 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -2503,13 +2503,37 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     # Non-indexed sizes are static above, so symbolic output sizes belong to
     # the broadcast index shape, which leads the gather result. Read them at
     # runtime instead of adding specialization guards through int(SymInt).
+    # Contiguous indexed axes keep the broadcast dimensions in place in ATen.
     leading_dims = axes[0] if axes == list(range(axes[0], axes[-1] + 1)) else 0
     out_shape = emit_shape(P, n, gather_slot, dim_offset=-leading_dims)
 
+    reshape_slot = gather_slot
+    broadcast_ndim = len(out_meta.shape) - x_ndim + len(axes)
+    broadcast_shape = out_meta.shape[leading_dims : leading_dims + broadcast_ndim]
+    # Moving singleton blocks does not change element order; keep those
+    # lowerings reshape-only, without extra instructions or tensor slots.
+    if any(size > 1 for size in slice_sizes[:leading_dims]) and any(
+        not isinstance(size, int) or size > 1 for size in broadcast_shape
+    ):
+        _, reshape_slot = P.make_tmp_slot()
+        P.emit(
+            TransposeNode(
+                x=P.slot_to_tid(gather_slot),
+                out=P.slot_to_tid(reshape_slot),
+                perm=(
+                    list(range(broadcast_ndim, broadcast_ndim + leading_dims))
+                    + list(range(broadcast_ndim))
+                    + list(
+                        range(broadcast_ndim + leading_dims, broadcast_ndim + x_ndim)
+                    )
+                ),
+            )
+        )
+
     out = P.make_or_get_slot(n)
     P.emit(
         ReshapeNode(
-            x=P.slot_to_tid(gather_slot),
+            x=P.slot_to_tid(reshape_slot),
             out=P.slot_to_tid(out),
             shape=out_shape,
         )
diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py
index ad460d49cf6..b4ee6f22ba1 100644
--- a/backends/mlx/test/test_ops.py
+++ b/backends/mlx/test/test_ops.py
@@ -4478,6 +4478,12 @@ class AdvancedIndexTest(OpTestCase):
     name = "advanced_index"
     rtol = 1e-4
     atol = 1e-4
+    expected_node_counts = {
+        "GatherNode": 1,
+        "ReshapeNode": 1,
+        "TransposeNode": 0,
+        "SymSizeNode": 0,
+    }
 
     def __init__(
         self,
@@ -4519,12 +4525,12 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
 
 
 class DynamicAdvancedIndexModel(nn.Module):
-    def __init__(self, leading_singleton: bool):
+    def __init__(self, has_leading_dim: bool):
         super().__init__()
-        self.leading_singleton = leading_singleton
+        self.has_leading_dim = has_leading_dim
 
     def forward(self, x, rows, columns):
-        if self.leading_singleton:
+        if self.has_leading_dim:
             return x[:, rows, columns]
         return x[rows, columns]
 
@@ -4540,16 +4546,22 @@ class DynamicAdvancedIndexTest(OpTestCase):
     def __init__(
         self,
         trailing_dim: bool = False,
-        leading_singleton: bool = False,
+        leading_dim: Optional[int] = None,
         test_rows: int = 4,
         test_columns: int = 5,
     ):
         self.trailing_dim = trailing_dim
-        self.leading_singleton = leading_singleton
+        self.leading_dim = leading_dim
+        self.expected_node_counts = {
+            "GatherNode": 1,
+            "ReshapeNode": 1,
+            "TransposeNode": int(leading_dim is not None and leading_dim > 1),
+            "SymSizeNode": 2,
+        }
         self.test_rows = test_rows
         self.test_columns = test_columns
         self.name = (
-            f"dynamic_advanced_index_tail{trailing_dim}_leading{leading_singleton}"
+            f"dynamic_advanced_index_tail{trailing_dim}_leading{leading_dim}"
             f"_runtime{test_rows}x{test_columns}"
         )
 
@@ -4558,21 +4570,27 @@ def get_test_configs(cls) -> List["DynamicAdvancedIndexTest"]:
         return [
             cls(
                 trailing_dim=trailing,
-                leading_singleton=leading,
+                leading_dim=leading,
                 test_rows=rows,
                 test_columns=columns,
             )
-            for trailing, leading in ((False, False), (True, False), (True, True))
+            for trailing, leading in (
+                (False, None),
+                (True, None),
+                (True, 1),
+                (False, 2),
+                (True, 2),
+            )
             for rows, columns in ((4, 5), (3, 2))
         ]
 
     def create_model(self) -> nn.Module:
-        return DynamicAdvancedIndexModel(self.leading_singleton)
+        return DynamicAdvancedIndexModel(self.leading_dim is not None)
 
     def _inputs(self, rows, columns):
         shape = (6, 7) + ((4,) if self.trailing_dim else ())
-        if self.leading_singleton:
-            shape = (1,) + shape
+        if self.leading_dim is not None:
+            shape = (self.leading_dim,) + shape
         return (
             torch.randn(shape),
             torch.arange(rows).reshape(-1, 1),
@@ -4593,6 +4611,58 @@ def get_dynamic_shapes(self) -> Optional[Dict]:
         }
 
 
+class AdvancedIndexLayoutModel(nn.Module):
+    def __init__(self, axes: Tuple[int, int]):
+        super().__init__()
+        self.axes = axes
+
+    def forward(self, x, rows, columns):
+        indices = [None] * x.ndim
+        indices[self.axes[0]] = rows
+        indices[self.axes[1]] = columns
+        return torch.ops.aten.index.Tensor(x, indices)
+
+
+@register_test
+class AdvancedIndexLayoutTest(OpTestCase):
+    """Preserve index ordering without transposing singleton or separated blocks."""
+
+    name = "advanced_index_layout"
+
+    def __init__(self, input_shape, axes, index_shape, expected_transposes):
+        self.input_shape = input_shape
+        self.axes = axes
+        self.index_shape = index_shape
+        self.name = f"advanced_index_layout_{input_shape}_{axes}_{index_shape}"
+        self.expected_node_counts = {
+            "GatherNode": 1,
+            "ReshapeNode": 1,
+            "TransposeNode": expected_transposes,
+            "SymSizeNode": 0,
+        }
+
+    @classmethod
+    def get_test_configs(cls) -> List["AdvancedIndexLayoutTest"]:
+        return [
+            cls((2, 6, 7), (1, 2), (2, 3), 1),
+            cls((2, 3, 6, 7, 4), (2, 3), (2, 3), 1),
+            cls((1, 1, 6, 7), (2, 3), (2, 3), 0),
+            cls((2, 6, 7), (1, 2), (), 0),
+            cls((2, 6, 7), (1, 2), (1, 1), 0),
+            cls((2, 6, 4, 7, 3), (1, 3), (2, 3), 0),
+        ]
+
+    def create_model(self) -> nn.Module:
+        return AdvancedIndexLayoutModel(self.axes)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (
+            torch.randn(self.input_shape),
+            torch.randint(self.input_shape[self.axes[0]], self.index_shape),
+            torch.randint(self.input_shape[self.axes[1]], self.index_shape),
+        )
+
+
 class IndexUpdateModel(nn.Module):
     """Model that performs index_copy on a mutable buffer.
 

From ef49f0ba8d6a4c222f3d60e667415c4ce88750bf Mon Sep 17 00:00:00 2001
From: Scott Roy <scroy@meta.com>
Date: Mon, 21 Sep 2026 15:54:22 -0700
Subject: [PATCH 3/4] up

---
 backends/mlx/passes.py                |  25 +++-
 backends/mlx/runtime/MLXInterpreter.h |  12 ++
 backends/mlx/test/test_ops.py         | 159 ++++++++++++++++++++++++
 backends/mlx/test/test_passes.py      | 172 +++++++++++++++++++++++---
 4 files changed, 345 insertions(+), 23 deletions(-)

diff --git a/backends/mlx/passes.py b/backends/mlx/passes.py
index 9ad23771be1..1c3c8b5c7f2 100644
--- a/backends/mlx/passes.py
+++ b/backends/mlx/passes.py
@@ -216,8 +216,9 @@ class CollapseDtypeConversionPass(ExportPass):
 
     _to_copy(dtype=bf16)(_to_copy(dtype=f32)(x)) → _to_copy(dtype=bf16)(x)
 
-    Only the final dtype matters. Only collapses when both nodes are pure dtype
-    conversions (no device/layout/memory_format changes).
+    Only collapse when the intermediate cast preserves every source value.
+    Narrowing or cross-kind casts may round, truncate, or overflow and must stay.
+    Both nodes must be pure dtype conversions (no device/layout/memory_format changes).
     """
 
     def call(self, graph_module: GraphModule) -> PassResult:
@@ -244,8 +245,26 @@ def call(self, graph_module: GraphModule) -> PassResult:
             if not _is_pure_dtype_cast(node_kw) or not _is_pure_dtype_cast(parent_kw):
                 continue
 
+            source = parent.args[0]
+            source_val = source.meta.get("val") if isinstance(source, Node) else None
+            if source_val is None:
+                continue
+            source_dtype = source_val.dtype
+            intermediate_dtype = parent_kw["dtype"]
+            if source_dtype != intermediate_dtype and (
+                source_dtype,
+                intermediate_dtype,
+            ) not in {
+                (torch.float16, torch.float32),
+                (torch.bfloat16, torch.float32),
+                (torch.float16, torch.float64),
+                (torch.bfloat16, torch.float64),
+                (torch.float32, torch.float64),
+            }:
+                continue
+
             # Rewrite: to_copy(to_copy(x, dtype=d1), dtype=d2) → to_copy(x, dtype=d2)
-            node.args = (parent.args[0],)
+            node.args = (source,)
             graph.erase_node(parent)
             modified = True
 
diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h
index 498c0e34f7c..1a3d03ba28b 100644
--- a/backends/mlx/runtime/MLXInterpreter.h
+++ b/backends/mlx/runtime/MLXInterpreter.h
@@ -12,6 +12,7 @@
 #include "MLXExecutor.h"
 
 #include <algorithm>
+#include <limits>
 #include <vector>
 
 #include <mlx/array.h>
@@ -302,6 +303,17 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) {
       sinks,
       false,
       s);
+  if (n.mask) {
+    const auto& M = st.const_tensor_ref(*n.mask);
+    array allowed = M.dtype() == bool_
+        ? M
+        : not_equal(
+              M, array(-std::numeric_limits<float>::infinity(), M.dtype()), s);
+    array has_key = M.ndim() == 0 ? allowed : any(allowed, -1, true, s);
+    // MLX can return nonzero values or NaNs for empty rows. Select, rather
+    // than multiply, to preserve PyTorch's zero-row semantics in both cases.
+    out = where(has_key, out, array(0, out.dtype()), s);
+  }
   st.set_tensor(n.out, std::move(out));
 }
 
diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py
index b4ee6f22ba1..89c0d77e56e 100644
--- a/backends/mlx/test/test_ops.py
+++ b/backends/mlx/test/test_ops.py
@@ -6061,6 +6061,69 @@ def create_model(self) -> nn.Module:
         return OnesModel(self.shape, self.dtype)
 
 
+class CastChainModel(nn.Module):
+    def __init__(self, intermediate_dtype: torch.dtype, output_dtype: torch.dtype):
+        super().__init__()
+        self.intermediate_dtype = intermediate_dtype
+        self.output_dtype = output_dtype
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Keep a delegated op even if the cast chain is incorrectly removed.
+        return x.to(self.intermediate_dtype).to(self.output_dtype) + 1
+
+
+@register_test
+class CastChainTest(OpTestCase):
+    """Default passes must preserve rounding and truncation in cast chains."""
+
+    name = "cast_chain"
+    rtol = 0
+    atol = 0
+
+    def __init__(
+        self,
+        intermediate_dtype: torch.dtype = torch.int32,
+        source_dtype: torch.dtype = torch.float32,
+        dynamic: bool = False,
+    ):
+        self.intermediate_dtype = intermediate_dtype
+        self.source_dtype = source_dtype
+        self.dynamic = dynamic
+        self.name = f"cast_chain_{source_dtype}_{intermediate_dtype}_{dynamic}"
+
+    @classmethod
+    def get_test_configs(cls) -> List["CastChainTest"]:
+        return [
+            cls(dtype, dynamic=dynamic)
+            for dtype in (torch.int32, torch.float16, torch.bfloat16)
+            for dynamic in (False, True)
+        ] + [
+            cls(torch.float32, source_dtype=dtype)
+            for dtype in (torch.float16, torch.bfloat16)
+        ]
+
+    def create_model(self) -> nn.Module:
+        return CastChainModel(self.intermediate_dtype, self.source_dtype)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (torch.tensor([1.9, -2.9, 3.7, 0.1], dtype=self.source_dtype),)
+
+    def create_test_inputs(self) -> Tuple[torch.Tensor, ...]:
+        x = self.create_inputs()[0]
+        return (torch.cat((x, x)) if self.dynamic else x,)
+
+    def get_dynamic_shapes(self) -> Optional[Dict]:
+        return {"x": {0: Dim("length", min=2, max=16)}} if self.dynamic else None
+
+    def get_edge_compile_config(self) -> Optional[exir.EdgeCompileConfig]:
+        return exir.EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True)
+
+    def get_transform_passes(self) -> Optional[list]:
+        from executorch.backends.mlx.passes import get_default_passes
+
+        return get_default_passes()
+
+
 class ToDtypeModel(nn.Module):
     def __init__(self, target_dtype: torch.dtype):
         super().__init__()
@@ -6434,6 +6497,102 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (q, k, v)
 
 
+class MaskedRowsSDPAModel(nn.Module):
+    def __init__(self, custom: bool):
+        super().__init__()
+        self.custom = custom
+
+    def forward(self, q, k, v, mask):
+        if self.custom:
+            return torch.ops.mlx.custom_sdpa(
+                q, k, v, start_pos=0, attn_mask=mask, is_causal=False
+            )
+        return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask)
+
+
+@register_test
+class SDPAMaskedRowsTest(OpTestCase):
+    """Empty rows must be zero, without zeroing finite biases or partial rows."""
+
+    name = "sdpa_masked_rows"
+    rtol = 0
+    atol = 0
+    expected_node_counts = {"SdpaNode": 1}
+
+    def __init__(
+        self,
+        mask_kind: str = "bool",
+        dtype: torch.dtype = torch.float32,
+        head_dim: int = 8,
+        seq_len: int = 4,
+        custom: bool = False,
+        per_batch: bool = False,
+    ):
+        self.mask_kind = mask_kind
+        self.dtype = dtype
+        self.head_dim = head_dim
+        self.seq_len = seq_len
+        self.custom = custom
+        self.per_batch = per_batch
+        self.name = (
+            f"sdpa_masked_rows_{mask_kind}_{dtype}_d{head_dim}_s{seq_len}"
+            f"_custom{custom}_batch{per_batch}"
+        )
+
+    @classmethod
+    def get_test_configs(cls) -> List["SDPAMaskedRowsTest"]:
+        return (
+            [
+                cls(mask_kind, dtype, head_dim, seq_len)
+                for mask_kind in ("bool", "additive")
+                for dtype in (torch.float32, torch.float16, torch.bfloat16)
+                for head_dim, seq_len in ((8, 4), (64, 4), (64, 16))
+            ]
+            + [
+                cls("finite", dtype)
+                for dtype in (torch.float32, torch.float16, torch.bfloat16)
+            ]
+            + [cls(mask_kind, custom=True) for mask_kind in ("bool", "additive")]
+            + [cls(mask_kind, per_batch=True) for mask_kind in ("bool", "additive")]
+        )
+
+    def create_model(self) -> nn.Module:
+        return MaskedRowsSDPAModel(self.custom)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        shape = (2, 2, self.seq_len, self.head_dim)
+        q = torch.zeros(shape, dtype=self.dtype)
+        k = torch.zeros_like(q)
+        v = (
+            torch.arange(1, self.seq_len + 1, dtype=self.dtype)
+            .view(1, 1, -1, 1)
+            .expand(shape)
+            .contiguous()
+        )
+        allowed = torch.ones(self.seq_len, self.seq_len, dtype=torch.bool)
+        allowed[0] = False
+        allowed[2, 1:] = False
+        if self.per_batch:
+            allowed = allowed.expand(2, 1, -1, -1).clone()
+            allowed[1, :, 0] = True
+        if self.mask_kind == "bool":
+            mask = allowed
+        else:
+            masked_value = (
+                torch.finfo(self.dtype).min
+                if self.mask_kind == "finite"
+                else float("-inf")
+            )
+            mask = torch.zeros(allowed.shape, dtype=self.dtype)
+            mask.masked_fill_(~allowed, masked_value)
+        return q, k, v, mask
+
+    def get_transform_passes(self) -> Optional[list]:
+        from executorch.backends.mlx.passes import get_default_passes
+
+        return get_default_passes()
+
+
 @register_test
 class SDPARank3Test(OpTestCase):
     """Attention on rank-3 tensors, which PyTorch accepts and the fused kernel does not.
diff --git a/backends/mlx/test/test_passes.py b/backends/mlx/test/test_passes.py
index dc4490d37b7..0b29b342e13 100644
--- a/backends/mlx/test/test_passes.py
+++ b/backends/mlx/test/test_passes.py
@@ -335,28 +335,138 @@ def forward(self, x):
 
 class TestCollapseDtypeConversionPass(unittest.TestCase):
 
-    def test_consecutive_casts_collapsed(self):
-        """_to_copy(f32→bf16→f16) → _to_copy(f32→f16)."""
+    def test_lossy_consecutive_casts_kept(self):
+        class M(nn.Module):
+            def __init__(self, intermediate, output):
+                super().__init__()
+                self.intermediate = intermediate
+                self.output = output
 
+            def forward(self, x):
+                return x.to(self.intermediate).to(self.output)
+
+        floats = torch.tensor([-1.003, -0.9, 0.9, 1.003], dtype=torch.float32)
+        cases = (
+            (floats, torch.int32, torch.float32),
+            (floats, torch.bfloat16, torch.float16),
+            (floats, torch.float16, torch.float32),
+            (
+                torch.tensor([2**24 + 1, 2**24 + 3, -(2**24 + 1)]),
+                torch.float32,
+                torch.int64,
+            ),
+        )
+        target = exir_ops.edge.aten._to_copy.default
+        for x, intermediate, output in cases:
+            with self.subTest(source=x.dtype, intermediate=intermediate, output=output):
+                model = M(intermediate, output)
+                expected = model(x)
+                self.assertFalse(torch.equal(expected, x.to(output)))
+                gm = _to_edge_gm(model, (x,))
+                self.assertEqual(_count_ops(gm, target), 2)
+
+                result = CollapseDtypeConversionPass()(gm)
+
+                self.assertFalse(result.modified)
+                self.assertEqual(_count_ops(result.graph_module, target), 2)
+                result.graph_module.recompile()
+                torch.testing.assert_close(
+                    result.graph_module(x)[0], expected, rtol=0, atol=0
+                )
+
+    def test_lossless_widening_casts_collapsed(self):
         class M(nn.Module):
+            def __init__(self, intermediate, output):
+                super().__init__()
+                self.intermediate = intermediate
+                self.output = output
+
             def forward(self, x):
-                return x.to(torch.bfloat16).to(torch.float16)
+                return x.to(self.intermediate).to(self.output)
+
+        cases = (
+            (torch.float16, torch.float32, torch.bfloat16),
+            (torch.bfloat16, torch.float32, torch.float16),
+            (torch.float16, torch.float64, torch.bfloat16),
+            (torch.bfloat16, torch.float64, torch.float16),
+            (torch.float32, torch.float64, torch.float16),
+        )
+        target = exir_ops.edge.aten._to_copy.default
+        for source, intermediate, output in cases:
+            with self.subTest(source=source, intermediate=intermediate):
+                x = torch.tensor([-1.003, -0.9, 0.9, 1.003], dtype=source)
+                model = M(intermediate, output)
+                gm = _to_edge_gm(model, (x,))
+                self.assertEqual(_count_ops(gm, target), 2)
+                source_node = _find_nodes(gm, target)[0].args[0]
+
+                result = CollapseDtypeConversionPass()(gm)
+
+                self.assertTrue(result.modified)
+                nodes = _find_nodes(result.graph_module, target)
+                self.assertEqual(len(nodes), 1)
+                self.assertIs(nodes[0].args[0], source_node)
+                self.assertEqual(nodes[0].kwargs["dtype"], output)
+                result.graph_module.recompile()
+                torch.testing.assert_close(
+                    result.graph_module(x)[0], model(x), rtol=0, atol=0
+                )
 
-        gm = _to_edge_gm(M(), (torch.randn(4, 4),))
+    def test_missing_source_metadata_not_collapsed(self):
+        class M(nn.Module):
+            def forward(self, x):
+                return x.to(torch.float32).to(torch.bfloat16)
+
+        x = torch.tensor([1.003, -0.9], dtype=torch.float16)
+        gm = _to_edge_gm(M(), (x,))
         target = exir_ops.edge.aten._to_copy.default
-        before = _count_ops(gm, target)
+        nodes = _find_nodes(gm, target)
+        self.assertEqual(len(nodes), 2)
+        del nodes[0].args[0].meta["val"]
+
+        result = CollapseDtypeConversionPass()(gm)
+
+        self.assertFalse(result.modified)
+        self.assertEqual(_count_ops(result.graph_module, target), 2)
+        torch.testing.assert_close(result.graph_module(x)[0], M()(x), rtol=0, atol=0)
+
+    def test_multi_user_parent_not_collapsed(self):
+        class M(nn.Module):
+            def forward(self, x):
+                y = x.to(torch.float32)
+                return y, y.to(torch.bfloat16)
 
-        if before < 2:
-            self.skipTest("Export optimized away double cast")
+        x = torch.tensor([1.003, -0.9], dtype=torch.float16)
+        gm = _to_edge_gm(M(), (x,))
+        target = exir_ops.edge.aten._to_copy.default
+        nodes = _find_nodes(gm, target)
+        self.assertEqual(len(nodes), 2)
+        self.assertEqual(len(nodes[0].users), 2)
 
         result = CollapseDtypeConversionPass()(gm)
 
-        self.assertTrue(result.modified)
-        self.assertEqual(_count_ops(result.graph_module, target), 1)
+        self.assertFalse(result.modified)
+        self.assertEqual(_count_ops(result.graph_module, target), 2)
+        torch.testing.assert_close(result.graph_module(x), M()(x), rtol=0, atol=0)
 
-        # Remaining cast should be to float16
-        nodes = _find_nodes(result.graph_module, target)
-        self.assertEqual(nodes[0].kwargs.get("dtype"), torch.float16)
+    def test_non_pure_cast_not_collapsed(self):
+        class M(nn.Module):
+            def forward(self, x):
+                return x.to(torch.float32).to(torch.bfloat16)
+
+        target = exir_ops.edge.aten._to_copy.default
+        for cast_index in (0, 1):
+            with self.subTest(cast_index=cast_index):
+                gm = _to_edge_gm(M(), (torch.ones(2, dtype=torch.float16),))
+                nodes = _find_nodes(gm, target)
+                self.assertEqual(len(nodes), 2)
+                node = nodes[cast_index]
+                node.kwargs = dict(node.kwargs, memory_format=torch.contiguous_format)
+
+                result = CollapseDtypeConversionPass()(gm)
+
+                self.assertFalse(result.modified)
+                self.assertEqual(_count_ops(result.graph_module, target), 2)
 
     def test_single_cast_unchanged(self):
         class M(nn.Module):
@@ -422,24 +532,46 @@ def forward(self, x):
         self.assertTrue(result.modified)
         self.assertFalse(_has_op(result.graph_module, target))
 
-    def test_identity_dtype_cast_removed_after_collapse(self):
-        """Chain: f32→f16→f32 collapses to f32→f32, then RemoveNoOps removes it."""
-
+    def test_lossy_dtype_roundtrip_kept_after_collapse(self):
         class M(nn.Module):
             def forward(self, x):
                 return x.to(torch.float16).to(torch.float32)
 
-        gm = _to_edge_gm(M(), (torch.randn(4, 4),))
+        x = torch.tensor([-1.003, -0.9, 0.9, 1.003], dtype=torch.float32)
+        expected = M()(x)
+        self.assertFalse(torch.equal(expected, x))
+        gm = _to_edge_gm(M(), (x,))
         target = exir_ops.edge.aten._to_copy.default
+        self.assertEqual(_count_ops(gm, target), 2)
 
-        if _count_ops(gm, target) < 2:
-            self.skipTest("Export optimized away double cast")
+        collapsed = CollapseDtypeConversionPass()(gm)
+        result = RemoveNoOpsPass()(collapsed.graph_module)
 
-        CollapseDtypeConversionPass()(gm)
-        result = RemoveNoOpsPass()(gm)
+        self.assertFalse(collapsed.modified)
+        self.assertFalse(result.modified)
+        self.assertEqual(_count_ops(result.graph_module, target), 2)
+        result.graph_module.recompile()
+        torch.testing.assert_close(result.graph_module(x)[0], expected, rtol=0, atol=0)
+
+    def test_lossless_dtype_roundtrip_removed_after_collapse(self):
+        class M(nn.Module):
+            def forward(self, x):
+                return x.to(torch.float32).to(torch.float16)
+
+        x = torch.tensor([-1.003, -0.9, 0.9, 1.003], dtype=torch.float16)
+        gm = _to_edge_gm(M(), (x,))
+        target = exir_ops.edge.aten._to_copy.default
+        self.assertEqual(_count_ops(gm, target), 2)
+
+        collapsed = CollapseDtypeConversionPass()(gm)
+        self.assertTrue(collapsed.modified)
+        self.assertEqual(_count_ops(collapsed.graph_module, target), 1)
+        result = RemoveNoOpsPass()(collapsed.graph_module)
 
         self.assertTrue(result.modified)
         self.assertEqual(_count_ops(result.graph_module, target), 0)
+        result.graph_module.recompile()
+        torch.testing.assert_close(result.graph_module(x)[0], M()(x), rtol=0, atol=0)
 
     def test_to_copy_with_memory_format_not_removed(self):
         """_is_pure_dtype_cast rejects kwargs with non-None memory_format."""

From 05ea3e18d6a9de120b5cb634b3763f4c8737dd18 Mon Sep 17 00:00:00 2001
From: Scott Roy <scroy@meta.com>
Date: Mon, 21 Sep 2026 17:36:42 -0700
Subject: [PATCH 4/4] up

---
 backends/mlx/ops.py              |  1 +
 backends/mlx/passes.py           |  5 +++
 backends/mlx/test/test_passes.py | 70 ++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 8e6e7b6ea3d..89a792e8d60 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -2505,6 +2505,7 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     # runtime instead of adding specialization guards through int(SymInt).
     # Contiguous indexed axes keep the broadcast dimensions in place in ATen.
     leading_dims = axes[0] if axes == list(range(axes[0], axes[-1] + 1)) else 0
+    # Read pre-transpose sizes; the offset maps broadcast dims back to gather axes.
     out_shape = emit_shape(P, n, gather_slot, dim_offset=-leading_dims)
 
     reshape_slot = gather_slot
diff --git a/backends/mlx/passes.py b/backends/mlx/passes.py
index 1c3c8b5c7f2..5d42ee7d3df 100644
--- a/backends/mlx/passes.py
+++ b/backends/mlx/passes.py
@@ -255,6 +255,11 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 source_dtype,
                 intermediate_dtype,
             ) not in {
+                # Boolean values 0 and 1 are exact in each floating-point dtype.
+                (torch.bool, torch.float16),
+                (torch.bool, torch.bfloat16),
+                (torch.bool, torch.float32),
+                (torch.bool, torch.float64),
                 (torch.float16, torch.float32),
                 (torch.bfloat16, torch.float32),
                 (torch.float16, torch.float64),
diff --git a/backends/mlx/test/test_passes.py b/backends/mlx/test/test_passes.py
index 0b29b342e13..3d36cbc8e2e 100644
--- a/backends/mlx/test/test_passes.py
+++ b/backends/mlx/test/test_passes.py
@@ -348,6 +348,7 @@ def forward(self, x):
         floats = torch.tensor([-1.003, -0.9, 0.9, 1.003], dtype=torch.float32)
         cases = (
             (floats, torch.int32, torch.float32),
+            (floats, torch.bool, torch.float32),
             (floats, torch.bfloat16, torch.float16),
             (floats, torch.float16, torch.float32),
             (
@@ -412,6 +413,46 @@ def forward(self, x):
                     result.graph_module(x)[0], model(x), rtol=0, atol=0
                 )
 
+    def test_boolean_to_float_casts_collapsed(self):
+        class M(nn.Module):
+            def __init__(self, intermediate, output):
+                super().__init__()
+                self.intermediate = intermediate
+                self.output = output
+
+            def forward(self, x):
+                return x.to(self.intermediate).to(self.output)
+
+        x = torch.tensor([False, True])
+        target = exir_ops.edge.aten._to_copy.default
+        for intermediate in (
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+            torch.float64,
+        ):
+            for output in (torch.float16, torch.float32, torch.int32):
+                if intermediate == output:
+                    continue
+                with self.subTest(intermediate=intermediate, output=output):
+                    model = M(intermediate, output)
+                    gm = _to_edge_gm(model, (x,))
+                    nodes = _find_nodes(gm, target)
+                    self.assertEqual(len(nodes), 2)
+                    source_node = nodes[0].args[0]
+
+                    result = CollapseDtypeConversionPass()(gm)
+
+                    self.assertTrue(result.modified)
+                    nodes = _find_nodes(result.graph_module, target)
+                    self.assertEqual(len(nodes), 1)
+                    self.assertIs(nodes[0].args[0], source_node)
+                    self.assertEqual(nodes[0].kwargs["dtype"], output)
+                    result.graph_module.recompile()
+                    torch.testing.assert_close(
+                        result.graph_module(x)[0], model(x), rtol=0, atol=0
+                    )
+
     def test_missing_source_metadata_not_collapsed(self):
         class M(nn.Module):
             def forward(self, x):
@@ -777,6 +818,35 @@ def forward(self, x):
         torch.testing.assert_close(result.module()(*inputs), model(*inputs))
         MLXProgramBuilder(result).build()
 
+    def test_boolean_cache_reset_uses_one_cast(self):
+        class M(nn.Module):
+            def forward(self, input_pos):
+                reset = (input_pos[0] == 0).to(torch.bfloat16).to(torch.float32)
+                return 1.0 - reset
+
+        model = M()
+        result = (
+            _to_edge(model, (torch.tensor([0]),))
+            .transform(get_default_passes())
+            .exported_program()
+        )
+        result.graph.lint()
+        for position in (0, 7):
+            with self.subTest(position=position):
+                input_pos = torch.tensor([position])
+                torch.testing.assert_close(
+                    result.module()(input_pos), model(input_pos), rtol=0, atol=0
+                )
+        built = MLXProgramBuilder(result).build()
+        self.assertEqual(
+            sum(
+                type(instr.op).__name__ == "AsTypeNode"
+                for chain in built.instruction_chains
+                for instr in chain.instructions
+            ),
+            1,
+        )
+
     def test_correctness_after_all_passes(self):
         """Output values should be preserved after running all passes."""