Add MLX op handler for aten.isinf (#18988)

Good day

## Summary

This PR adds a decomposed MLX op handler for `aten.isinf` to the
pytorch/executorch project.

### Motivation

The MLX delegate converts PyTorch aten ops into MLX graph nodes during
export. When an aten op has no handler, it falls back to CPU execution,
breaking the GPU acceleration pipeline. Adding a handler for
`aten.isinf` enables it to run on Metal GPU via MLX.

### Implementation

The handler uses a decomposed approach:
```python
isinf(x) = abs(x) == inf
```

This uses existing `AbsNode` and `EqualNode` which are already
supported, avoiding the need for a dedicated MLX isinf op.

### Changes

- **backends/mlx/ops.py**: Added `_isinf_handler` function registered
for `torch.ops.aten.isinf.default`
- **backends/mlx/test/test_ops.py**: Added `isinf` to `_UNARY_OP_TESTS`
with standard test configuration

### Testing

The handler can be tested with:
```bash
python -m executorch.backends.mlx.test.run_all_tests -k isinf
```

Thank you for your attention. If there are any issues or suggestions,
please leave a comment and I will address them promptly.

Warmly,
RoomWithOutRoof

Co-authored-by: Scott Roy <161522778+metascroy@users.noreply.github.com>

Author: Jah-yee

backends/mlx/ops.py

@@ -666,6 +666,41 @@ def _isnan_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+@REGISTRY.register(target=[torch.ops.aten.isinf.default])
+def _isinf_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Handle aten.isinf - check for infinite values element-wise.
+
+    isinf(x) is equivalent to abs(x) == inf.
+    """
+    args = P.args(n)
+    require_args(args, 1, 1, "aten.isinf")
+    require_kwargs(P.kwargs(n), set(), "aten.isinf")
+    x = args[0]
+
+    # Create abs(x)
+    _, abs_tmp = P.make_tmp_slot()
+    P.emit(
+        AbsNode(
+            x=P.slot_to_tid(x),
+            out=P.slot_to_tid(abs_tmp),
+        )
+    )
+
+    # Create inf constant (float32; EqualNode handles type promotion to match input dtype)
+    inf_slot = emit_lifted_constant(P, float("inf"), torch.float32)
+
+    # Compare abs(x) == inf
+    out = P.make_or_get_slot(n)
+    P.emit(
+        EqualNode(
+            a=P.slot_to_tid(abs_tmp),
+            b=P.slot_to_tid(inf_slot),
+            out=P.slot_to_tid(out),
+        )
+    )
+    return out
+
+
 _BINARY_OPS: List[Tuple[List[Any], Any, str, bool]] = [
     (
         [torch.ops.aten.mul.Tensor, torch.ops.aten.mul.Scalar],

backends/mlx/test/test_ops.py

@@ -4737,6 +4737,21 @@ def fn(shape, dtype):
     return fn
 
 
+def _inf_input_fn():
+    """Return a callable(shape, dtype) that generates inputs with some inf values."""
+
+    def fn(shape, dtype):
+        x = torch.randn(shape, dtype=dtype)
+        # Insert ~20% +inf and ~10% -inf using non-overlapping masks
+        mask_pos = torch.rand(shape) > 0.8  # ~20% -> +inf
+        mask_neg = (~mask_pos) & (torch.rand(shape) > 0.9)  # ~10% of remaining -> -inf
+        x[mask_pos] = float("inf")
+        x[mask_neg] = float("-inf")
+        return (x,)
+
+    return fn
+
+
 # Standard shape and dtype configs used by unary tests.
 _SHAPES_3 = [(16,), (4, 4), (2, 3, 4)]
 _SHAPES_2 = [(16,), (4, 4)]
@@ -4830,6 +4845,7 @@ def create_model(self) -> nn.Module:
     {"op_name": "logical_not","op_fn": torch.logical_not, "shapes": [(2, 3, 4), (10,), (4, 8)], "dtypes": [torch.bool], "input_fn": _bool_input_fn()},
     {"op_name": "bitwise_not_int", "op_fn": torch.bitwise_not, "shapes": _SHAPES_3, "dtypes": [torch.int32, torch.int64], "input_fn": _int_input_fn()},
     {"op_name": "isnan",      "op_fn": torch.isnan,      "shapes": _SHAPES_3, "dtypes": [torch.float32, torch.float16, torch.bfloat16], "input_fn": _nan_input_fn()},
+    {"op_name": "isinf",      "op_fn": torch.isinf,      "shapes": _SHAPES_3, "dtypes": [torch.float32, torch.float16, torch.bfloat16], "input_fn": _inf_input_fn()},
     # activations
     {"op_name": "relu",    "op_fn": torch.relu,    "shapes": [(2, 3, 4), (10,), (4, 8), (2, 8, 16), (1, 128, 64)], "dtypes": [torch.float32], "input_fn": _input_fn(scale=2, offset=-1)},
     {"op_name": "sigmoid", "op_fn": torch.sigmoid, "shapes": [(2, 3, 4), (10,), (4, 8), (2, 8, 16), (1, 1, 128)],  "dtypes": [torch.float32], "input_fn": _input_fn(scale=2)},