Fix MLX advanced indexing and support boolean masks

pytensor/link/mlx/dispatch/subtensor.py

@@ -26,8 +26,27 @@ def subtensor(x, *ilists):
 
 @mlx_funcify.register(AdvancedSubtensor)
 def mlx_funcify_AdvancedSubtensor(op, node, **kwargs):
-    def advanced_subtensor(x, *ilists):
-        indices = indices_from_subtensor(ilists, op.idx_list)
+    # MLX slices reject array-typed bounds, so every flat index position a slice
+    # uses as a bound is coerced to a Python int; array-valued indices are left
+    # alone.
+    bound_positions = {
+        bound
+        for idx_entry in op.idx_list
+        if isinstance(idx_entry, slice)
+        for bound in (idx_entry.start, idx_entry.stop, idx_entry.step)
+        if bound is not None
+    }
+
+    def advanced_subtensor(
+        x, *ilists, idx_list=op.idx_list, bound_positions=bound_positions
+    ):
+        indices = indices_from_subtensor(
+            tuple(
+                int(index) if position in bound_positions else index
+                for position, index in enumerate(ilists)
+            ),
+            idx_list,
+        )
         if len(indices) == 1:
             indices = indices[0]
 
@@ -96,9 +115,40 @@ def mlx_fn(x, indices, y):
         def mlx_fn(x, indices, y):
             return x.at[indices].add(y)
 
-    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn):
+    # MLX slices reject array-typed bounds, so every flat index position a slice
+    # uses as a bound is coerced to a Python int; array-valued indices are left
+    # alone.
+    bound_positions = {
+        bound
+        for idx_entry in op.idx_list
+        if isinstance(idx_entry, slice)
+        for bound in (idx_entry.start, idx_entry.stop, idx_entry.step)
+        if bound is not None
+    }
+
+    def advancedincsubtensor(
+        x,
+        y,
+        *ilist,
+        mlx_fn=mlx_fn,
+        idx_list=op.idx_list,
+        bound_positions=bound_positions,
+    ):
         op._check_runtime_broadcast_of_vector_index(node, x, y, ilist[0])
 
-        return mlx_fn(x, ilist, y)
+        # Slices and plain integers live in `idx_list`, not in `ilist`, and have
+        # to be spliced back in; without them `x[:, idx] = y` would scatter
+        # along the leading axis instead.
+        indices = indices_from_subtensor(
+            tuple(
+                int(index) if position in bound_positions else index
+                for position, index in enumerate(ilist)
+            ),
+            idx_list,
+        )
+        if len(indices) == 1:
+            indices = indices[0]
+
+        return mlx_fn(x, indices, y)
 
     return advancedincsubtensor

pytensor/link/mlx/dispatch/tensor_basic.py

@@ -11,6 +11,7 @@
     Eye,
     Join,
     MakeVector,
+    Nonzero,
     ScalarFromTensor,
     Split,
     TensorFromScalar,
@@ -206,6 +207,23 @@ def arange(*_args):
     return arange
 
 
+@mlx_funcify.register(Nonzero)
+def mlx_funcify_Nonzero(op, node, **kwargs):
+    ndim = node.inputs[0].type.ndim
+
+    def nonzero(a):
+        # How many entries are selected is data-dependent, so `a` has to be
+        # counted on the host. `MLXLinker` keeps graphs containing `Nonzero` out
+        # of `mx.compile`, which forbids the evaluation this needs.
+        indices = np.nonzero(np.asarray(a))
+        if ndim == 1:
+            return mx.array(indices[0])
+
+        return [mx.array(index) for index in indices]
+
+    return nonzero
+
+
 def _extract_static_dims(shape_inputs):
     static_dims = []
     for dim in shape_inputs:

pytensor/link/mlx/linker.py

@@ -1,6 +1,41 @@
+import warnings
+
 from pytensor.link.basic import JITLinker
 
 
+UNCOMPILABLE_SHAPE_WARNING = (
+    "This graph contains an `Op` whose output shape depends on its input "
+    "values, which `mx.compile` cannot trace, so the whole graph will run "
+    "uncompiled. To get compilation back, give the graph static shapes: a "
+    "boolean mask that is constant becomes integer indices, and a mask only "
+    "known at runtime can often be replaced by shape-preserving arithmetic, "
+    "e.g. `(x * mask).sum()` rather than `x[mask].sum()`."
+)
+
+
+def _has_data_dependent_shape(fgraph):
+    """Whether any `Op` in ``fgraph`` has an output shape that depends on data.
+
+    `mx.compile` traces a graph once per input *shape*, so an output whose shape
+    is only known from the values cannot be traced at all.
+    """
+    # Imported here because `pytensor.compile` imports this module while
+    # `pytensor.tensor` is still initializing.
+    from pytensor.graph.op import HasInnerGraph
+    from pytensor.tensor.basic import Nonzero
+
+    for node in fgraph.apply_nodes:
+        if isinstance(node.op, Nonzero):
+            return True
+
+        if isinstance(node.op, HasInnerGraph):
+            inner_fgraph = getattr(node.op, "fgraph", None)
+            if inner_fgraph is not None and _has_data_dependent_shape(inner_fgraph):
+                return True
+
+    return False
+
+
 class MLXLinker(JITLinker):
     """A `Linker` that JIT-compiles NumPy-based operations using Apple's MLX."""
 
@@ -44,7 +79,12 @@ def jit_compile(self, fn):
 
         from pytensor.link.mlx.dispatch import mlx_typify
 
-        if not self.use_compile:
+        use_compile = self.use_compile
+        if use_compile and _has_data_dependent_shape(self.fgraph):
+            warnings.warn(UNCOMPILABLE_SHAPE_WARNING, UserWarning, stacklevel=2)
+            use_compile = False
+
+        if not use_compile:
             # Skip compilation and just return the function with MLX typification
             def fn_no_compile(*inputs):
                 return fn(*(mlx_typify(inp) for inp in inputs))

pytensor/tensor/rewriting/jax.py

@@ -48,7 +48,10 @@ def boolean_indexing_set_or_inc(fgraph, node):
     "jax_boolean_indexing_set_or_inc",
     dfs_rewriter(boolean_indexing_set_or_inc),
     "jax",
-    position=100,
+    "mlx",
+    # Must run before `specialize`: `bool_idx_to_nonzero` lives there and rewrites
+    # the mask into integer indices, leaving nothing for this to match.
+    position=1.95,
 )
 
 
@@ -96,7 +99,12 @@ def boolean_indexing_sum(fgraph, node):
 
 
 optdb.register(
-    "jax_boolean_indexing_sum", dfs_rewriter(boolean_indexing_sum), "jax", position=100
+    "jax_boolean_indexing_sum",
+    dfs_rewriter(boolean_indexing_sum),
+    "jax",
+    "mlx",
+    # Before `specialize`, for the same reason as `boolean_indexing_set_or_inc`.
+    position=1.95,
 )
 
 

pytensor/tensor/rewriting/subtensor.py

@@ -2947,6 +2947,7 @@ def bool_idx_to_nonzero(fgraph, node):
     bool_idx_to_nonzero.__name__,
     bool_idx_to_nonzero,
     "numba",
+    "mlx",
     "shape_unsafe",  # It can mask invalid mask sizes
     use_db_name_as_tag=False,  # Not included if only "specialize" is requested
 )