Only take the single-bool-mask index path when the mask is the only index

coremltools/converters/mil/frontend/torch/ops.py

@@ -5739,16 +5739,17 @@ def _try_single_bool_index(x: Var, indices: List[Var], name: str) -> bool:
             The true value indicates whether the element should be selected among the masked axes
             The output c is a tensor with shape (2, N), where N is the number of elements of b satisfying condition > 0.1
         """
-        boolean_indices_axis = []
-        for i, index in enumerate(indices):
-            if index is not None and types.is_bool(index.dtype):
-                boolean_indices_axis.append(i)
+        non_none_indices_axis = [i for i, index in enumerate(indices) if index is not None]
+        boolean_indices_axis = [i for i in non_none_indices_axis if types.is_bool(indices[i].dtype)]
 
-        if len(boolean_indices_axis) == 1:
+        # This shortcut only holds when the mask is the sole index. With another index present,
+        # e.g. x[mask, j], torch pairs the mask's True positions up with j elementwise instead of
+        # selecting whole slices, which is what the general path below does.
+        if len(boolean_indices_axis) == 1 and len(non_none_indices_axis) == 1:
             # get the True element indices
             axis = boolean_indices_axis[0]
-            axes = list(range(axis, axis + index.rank))
             index = indices[axis]
+            axes = list(range(axis, axis + index.rank))
             index = mb.non_zero(x=index)
 
             # transpose the masked axes to the beginning

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -11372,6 +11372,33 @@ def forward(self, x, y):
                     minimum_deployment_target=minimum_deployment_target,
                 )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend",
+        itertools.product(compute_units, backends, frontends),
+    )
+    def test_index_bool_mask_with_another_index(self, compute_unit, backend, frontend):
+        """A bool mask paired with a second index selects elements, not whole slices."""
+        if frontend in TORCH_EXPORT_BASED_FRONTENDS:
+            pytest.xfail(
+                "torch.export cannot trace this model: the bool mask makes the intermediate "
+                "size data dependent, so export raises PendingUnbackedSymbolNotFound before "
+                "the converter is reached."
+            )
+
+        class IndexModel(torch.nn.Module):
+            def forward(self, x):
+                mask = torch.tensor([True, True])
+                j = torch.tensor([0, 2])
+                return x[mask, j]
+
+        self.run_compare_torch(
+            [(2, 3)],
+            IndexModel(),
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+        )
+
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, input_dtype, shape, minimum_deployment_target",
         itertools.product(