Arm backend: Add support for isinf and isnan

backends/arm/_passes/__init__.py

@@ -74,6 +74,7 @@
     DecomposeIndexTensorToGatherPass,
 )
 from .decompose_int_pow_pass import DecomposeIntPowPass  # noqa
+from .decompose_isinf_isnan_pass import DecomposeIsInfAndIsNanPass  # noqa
 from .decompose_large_stride_maxpool2d_pass import (  # noqa
     DecomposeLargeStrideMaxPool2dForU55Pass,
 )

backends/arm/_passes/arm_pass_manager.py

@@ -69,6 +69,7 @@
     DecomposeIndexSelectToGatherPass,
     DecomposeIndexTensorToGatherPass,
     DecomposeIntPowPass,
+    DecomposeIsInfAndIsNanPass,
     DecomposeLargeStrideMaxPool2dForU55Pass,
     DecomposeLayerNormPass,
     DecomposeLeakyReLUPass,
@@ -574,6 +575,7 @@ def _tosa_pipeline(
                 RemoveGetItemPass(),
                 FuseBatchNorm2dPass(exported_program),
                 DecomposeBatchNormNoStatsPass(),
+                DecomposeIsInfAndIsNanPass(),
                 DecomposeLogitPass(),
                 DecomposeMaskedFillPass(),
                 DecomposeRoundPass(),

backends/arm/_passes/decompose_isinf_isnan_pass.py

@@ -0,0 +1,48 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeIsInfAndIsNanPass(ArmOpTargetedPass):
+    """Decompose ``isinf`` and ``isnan`` into TOSA-supported operations."""
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+    edge_isinf = exir_ops.edge.aten.isinf.default
+    edge_isnan = exir_ops.edge.aten.isnan.default
+    target_ops = (edge_isinf, edge_isnan)
+    check_allowed_to_transform = True
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self.target_ops or not self.allowed_to_transform(meta):
+            return super().call_operator(op, args, kwargs, meta)
+
+        (x,) = args
+        abs_op = exir_ops.edge.aten.abs.default
+        eq_op = exir_ops.edge.aten.eq.Tensor
+        logical_not_op = exir_ops.edge.aten.logical_not.default
+        full_op = exir_ops.edge.aten.full.default
+
+        if op is self.edge_isnan:
+            equal = super().call_operator(eq_op, (x, x), {}, meta, updated=True)
+            return super().call_operator(
+                logical_not_op, (equal,), {}, meta, updated=True
+            )
+
+        absolute = super().call_operator(abs_op, (x,), {}, meta, updated=True)
+        infinity = super().call_operator(
+            full_op,
+            (x.data.shape, float("inf")),
+            {"dtype": x.data.dtype},
+            meta,
+            updated=True,
+        )
+        return super().call_operator(
+            eq_op, (absolute, infinity), {}, meta, updated=True
+        )

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -186,6 +186,8 @@
     exir_ops.edge.aten.expm1.default,
     exir_ops.edge.aten.log1p.default,
     exir_ops.edge.aten.log.default,
+    exir_ops.edge.aten.isnan.default,
+    exir_ops.edge.aten.isinf.default,
     exir_ops.edge.aten.linear.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,
     exir_ops.edge.aten.split_copy.Tensor,

backends/arm/operator_support/tosa_supported_operators.py

@@ -1151,7 +1151,7 @@ def is_node_supported(
 class CheckFPComparisonInputs(OperatorSupportBase):
     """Reject unsupported comparison inputs under the FP profile."""
 
-    target_ops = {
+    comparison_ops = {
         exir_ops.edge.aten.eq.Tensor,
         exir_ops.edge.aten.eq.Scalar,
         exir_ops.edge.aten.ne.Tensor,
@@ -1165,6 +1165,10 @@ class CheckFPComparisonInputs(OperatorSupportBase):
         exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.lt.Scalar,
     }
+    target_ops = comparison_ops | {
+        exir_ops.edge.aten.isinf.default,
+        exir_ops.edge.aten.isnan.default,
+    }
     supported_dtypes = {torch.float16, torch.float32, torch.bfloat16}
     castable_comparison_dtypes = {torch.int8, torch.int16}
 
@@ -1186,7 +1190,9 @@ def is_node_supported(
         if all(dtype in self.supported_dtypes for dtype in input_dtypes):
             return True
 
-        if all(dtype in self.castable_comparison_dtypes for dtype in input_dtypes):
+        if node.target in self.comparison_ops and all(
+            dtype in self.castable_comparison_dtypes for dtype in input_dtypes
+        ):
             return True
 
         unsupported_dtype = next(

backends/arm/test/ops/test_isinf.py

@@ -0,0 +1,79 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections.abc import Callable
+
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    OpNotSupportedPipeline,
+    TosaPipelineFP,
+    VgfPipeline,
+)
+
+aten_op = "torch.ops.aten.isinf.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten_isinf_default"
+
+
+class IsInf(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.isinf(x)
+
+
+test_data_suite = {
+    "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]),
+    "inf": lambda: torch.tensor([-float("inf"), 0.0, float("inf"), float("nan")]),
+    "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32),
+    "rank4": lambda: torch.tensor(
+        [[[[float("inf"), 0.0]]], [[[-float("inf"), float("nan")]]]]
+    ),
+}
+
+
+@common.parametrize(
+    "test_data",
+    {name: data for name, data in test_data_suite.items() if name != "integer"},
+)
+def test_isinf_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None:
+    TosaPipelineFP(
+        IsInf(),
+        (test_data(),),
+        aten_op,
+        exir_op,
+    ).run()
+
+
+def test_isinf_tosa_FP_falls_back_for_integer() -> None:
+    OpNotSupportedPipeline(
+        IsInf(),
+        (test_data_suite["integer"](),),
+        {exir_op: 1},
+        quantize=False,
+    ).run()
+
+
+def test_isinf_tosa_INT_falls_back() -> None:
+    test_data = (test_data_suite["inf"](),)
+    pipeline = OpNotSupportedPipeline(
+        IsInf(),
+        test_data,
+        {exir_op: 1},
+        quantize=True,
+    )
+    quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0]
+    quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)]
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_isinf_vgf_no_quant() -> None:
+    VgfPipeline(
+        IsInf(),
+        (test_data_suite["inf"](),),
+        aten_op,
+        exir_op,
+        quantize=False,
+    ).run()

backends/arm/test/ops/test_isnan.py

@@ -0,0 +1,77 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections.abc import Callable
+
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    OpNotSupportedPipeline,
+    TosaPipelineFP,
+    VgfPipeline,
+)
+
+aten_op = "torch.ops.aten.isnan.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten_isnan_default"
+
+
+class IsNan(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.isnan(x)
+
+
+test_data_suite = {
+    "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]),
+    "nan": lambda: torch.tensor([float("nan"), 0.0, float("inf")]),
+    "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32),
+    "rank4": lambda: torch.tensor([[[[float("nan"), 0.0]]], [[[float("inf"), -3.0]]]]),
+}
+
+
+@common.parametrize(
+    "test_data",
+    {name: data for name, data in test_data_suite.items() if name != "integer"},
+)
+def test_isnan_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None:
+    TosaPipelineFP(
+        IsNan(),
+        (test_data(),),
+        aten_op,
+        exir_op,
+    ).run()
+
+
+def test_isnan_tosa_FP_falls_back_for_integer() -> None:
+    OpNotSupportedPipeline(
+        IsNan(),
+        (test_data_suite["integer"](),),
+        {exir_op: 1},
+        quantize=False,
+    ).run()
+
+
+def test_isnan_tosa_INT_falls_back() -> None:
+    test_data = (test_data_suite["nan"](),)
+    pipeline = OpNotSupportedPipeline(
+        IsNan(),
+        test_data,
+        {exir_op: 1},
+        quantize=True,
+    )
+    quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0]
+    quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)]
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_isnan_vgf_no_quant() -> None:
+    VgfPipeline(
+        IsNan(),
+        (test_data_suite["nan"](),),
+        aten_op,
+        exir_op,
+        quantize=False,
+    ).run()

backends/arm/test/targets.bzl

@@ -19,6 +19,8 @@ def define_arm_tests():
         "ops/test_avg_pool2d.py",
         "ops/test_cat.py",
         "ops/test_conv2d.py",
+        "ops/test_isinf.py",
+        "ops/test_isnan.py",
         "ops/test_linear.py",
         "ops/test_log10.py",
         "ops/test_max_pool1d.py",

docs/source/backends/arm-vgf/VGF_op_support.md

@@ -6,7 +6,7 @@ This page lists VGF-supported PyTorch APIs and the dtype and quantization modes
 
 `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights.
 
-Total supported PyTorch APIs: **154**.
+Total supported PyTorch APIs: **156**.
 
 | PyTorch API | Support profile | DType | Quantization mode |
 | --- | --- | --- | --- |
@@ -69,6 +69,8 @@ Total supported PyTorch APIs: **154**.
 | `torch.gt` / `>` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.index_put_` | INT | `INT8` | 8x8 |
 | `torch.index_select` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `BOOL` | 8x8 |
+| `torch.isinf` | FP | `FP32` | - |
+| `torch.isnan` | FP | `FP32` | - |
 | `torch.layer_norm` | FP, INT | `FP32`, `INT8` | 8x8 |
 | `torch.le` / `<=` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.linspace` | FP, INT | `FP32`, `INT8` | 8x8 |