Refactor MLX subtensor dispatch and update test

Simplifies index normalization logic in MLX subtensor dispatch functions by separating basic and advanced indexing cases. Updates the advanced incsubtensor test to use vector array indices and a matching value shape for improved coverage.

Author: cetagostini

pytensor/link/mlx/dispatch/subtensor.py

@@ -16,7 +16,7 @@
 from pytensor.tensor.type_other import MakeSlice
 
 
-def normalize_indices_for_mlx(ilist, idx_list):
+def normalize_indices_for_mlx(indices):
     """Convert numpy integers to Python integers for MLX indexing.
 
     MLX requires index values to be Python int, not np.int64 or other NumPy types.
@@ -49,18 +49,19 @@ def normalize_element(element):
         else:
             return element
 
-    indices = indices_from_subtensor(ilist, idx_list)
     return tuple(normalize_element(idx) for idx in indices)
 
 
 @mlx_funcify.register(Subtensor)
 def mlx_funcify_Subtensor(op, node, **kwargs):
     """MLX implementation of Subtensor."""
-    idx_list = getattr(op, "idx_list", None)
+    idx_list = op.idx_list
 
     def subtensor(x, *ilists):
+        # Convert ilist to indices using idx_list (basic subtensor)
+        indices = indices_from_subtensor(ilists, idx_list)
         # Normalize indices to handle np.int64 and other NumPy types
-        indices = normalize_indices_for_mlx(ilists, idx_list)
+        indices = normalize_indices_for_mlx(indices)
         if len(indices) == 1:
             indices = indices[0]
 
@@ -73,11 +74,11 @@ def subtensor(x, *ilists):
 @mlx_funcify.register(AdvancedSubtensor1)
 def mlx_funcify_AdvancedSubtensor(op, node, **kwargs):
     """MLX implementation of AdvancedSubtensor."""
-    idx_list = getattr(op, "idx_list", None)
 
     def advanced_subtensor(x, *ilists):
         # Normalize indices to handle np.int64 and other NumPy types
-        indices = normalize_indices_for_mlx(ilists, idx_list)
+        # Advanced indexing doesn't use idx_list or indices_from_subtensor
+        indices = normalize_indices_for_mlx(ilists)
         if len(indices) == 1:
             indices = indices[0]
 
@@ -87,12 +88,11 @@ def advanced_subtensor(x, *ilists):
 
 
 @mlx_funcify.register(IncSubtensor)
-@mlx_funcify.register(AdvancedIncSubtensor1)
 def mlx_funcify_IncSubtensor(op, node, **kwargs):
     """MLX implementation of IncSubtensor."""
-    idx_list = getattr(op, "idx_list", None)
+    idx_list = op.idx_list
 
-    if getattr(op, "set_instead_of_inc", False):
+    if op.set_instead_of_inc:
 
         def mlx_fn(x, indices, y):
             if not op.inplace:
@@ -109,8 +109,10 @@ def mlx_fn(x, indices, y):
             return x
 
     def incsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=idx_list):
+        # Convert ilist to indices using idx_list (basic inc_subtensor)
+        indices = indices_from_subtensor(ilist, idx_list)
         # Normalize indices to handle np.int64 and other NumPy types
-        indices = normalize_indices_for_mlx(ilist, idx_list)
+        indices = normalize_indices_for_mlx(indices)
 
         if len(indices) == 1:
             indices = indices[0]
@@ -121,11 +123,11 @@ def incsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=idx_list):
 
 
 @mlx_funcify.register(AdvancedIncSubtensor)
+@mlx_funcify.register(AdvancedIncSubtensor1)
 def mlx_funcify_AdvancedIncSubtensor(op, node, **kwargs):
     """MLX implementation of AdvancedIncSubtensor."""
-    idx_list = getattr(op, "idx_list", None)
 
-    if getattr(op, "set_instead_of_inc", False):
+    if op.set_instead_of_inc:
 
         def mlx_fn(x, indices, y):
             if not op.inplace:
@@ -141,9 +143,10 @@ def mlx_fn(x, indices, y):
             x[indices] += y
             return x
 
-    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=idx_list):
+    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn):
         # Normalize indices to handle np.int64 and other NumPy types
-        indices = normalize_indices_for_mlx(ilist, idx_list)
+        # Advanced indexing doesn't use idx_list or indices_from_subtensor
+        indices = normalize_indices_for_mlx(ilist)
 
         # For advanced indexing, if we have a single tuple of indices, unwrap it
         if len(indices) == 1:

tests/link/mlx/test_subtensor.py

@@ -395,13 +395,16 @@ def test_mlx_advanced_incsubtensor_with_numpy_int64():
     x_np = np.arange(15, dtype=np.float32).reshape((5, 3))
     x_pt = pt.constant(x_np)
 
-    # Value to set/increment
+    # Value to set/increment - using 4 rows now for vector indexing
     y_pt = pt.as_tensor_variable(
-        np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
+        np.array(
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
+            dtype=np.float32,
+        )
     )
 
-    # Advanced indexing set with array indices
-    indices = [np.int64(0), np.int64(2)]
+    # Advanced indexing set with vector array indices
+    indices = np.array([0, 1, 2, 3], dtype=np.int64)
     out_pt = pt_subtensor.set_subtensor(x_pt[indices], y_pt)
     compare_mlx_and_py([], [out_pt], [])