Add MLX op handler for aten.isinf

Add decomposed handler for aten.isinf using AbsNode + EqualNode
with an inf constant. This enables isinf to run on Metal GPU via MLX
instead of falling back to CPU execution.

Also adds isinf to the unary op test suite with standard test inputs.

Author: Jah-yee

backends/mlx/ops.py

@@ -397,6 +397,46 @@ def normalize_reduction_dim(
 ]
 
 
+@REGISTRY.register(target=[torch.ops.aten.isinf.default])
+def _isinf_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Handle aten.isinf.default - check for infinite values element-wise.
+
+    isinf(x) is equivalent to abs(x) == inf.
+    This decomposition avoids a dedicated isinf MLX op by using AbsNode + EqualNode.
+    """
+    args = P.args(n)
+    require_args(args, 1, 1, "aten.isinf")
+    require_kwargs(P.kwargs(n), set(), "aten.isinf")
+    x = args[0]
+
+    # Get input dtype for the constant
+    x_meta = n.args[0].meta.get("val")
+    dtype = x_meta.dtype if x_meta is not None else torch.float32
+
+    # Create abs(x) using a temporary slot
+    _, 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
+    inf_slot = emit_lifted_constant(P, float("inf"), dtype)
+
+    # 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
+
+
 def _make_unary_handler(node_cls: Any, op_name: str):
     """Create a handler for a simple unary op: x → node_cls(x, out)."""
 

backends/mlx/test/test_ops.py

@@ -4095,6 +4095,7 @@ def create_model(self) -> nn.Module:
     {"op_name": "abs",        "op_fn": torch.abs},
     {"op_name": "neg",        "op_fn": torch.neg},
     {"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": "isinf",      "op_fn": torch.isinf,     "shapes": [(2, 3, 4), (10,), (4, 8)], "dtypes": [torch.float32], "input_fn": _input_fn(uniform=True)},  # isinf uses custom inputs via model wrapper
     # 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)},