Add MLX op handler for aten.isinf

[Jah-yee]

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:

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:

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

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18988

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI runner label rename: rebase PRs using old linux.rocm.gpu.gfx950.* labels

⚠️ 13 Awaiting Approval

As of commit f811554 with merge base 490ec5c ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• MLX (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[nil-is-all]

@claude review

[claude]

Claude finished @nil-is-all's task in 1m 31s —— View job

---

Reviewing PR #18988: Add MLX op handler for aten.isinf

• Gather context and read changed files
• Review backends/mlx/ops.py changes
• Review backends/mlx/test/test_ops.py changes
• Post review feedback

---

Review Summary

The decomposition approach (isinf(x) = abs(x) == inf) is sound and follows the same pattern as the existing _isnan_handler. The implementation correctly uses AbsNode and EqualNode, and the test coverage is appropriate. A few items to address:

---

Issues

1. Formatting: # activations comment merged onto the isinf line — backends/mlx/test/test_ops.py:4130

The isinf entry was inserted before the # activations comment, but the comment ended up appended to the end of the isinf line instead of staying on its own line:

# Before (line 4129-4131 in original):
    {"op_name": "isnan", ...},
    # activations
    {"op_name": "relu", ...},

# After (current):
    {"op_name": "isnan", ...},
    {"op_name": "isinf", ...},    # activations
    {"op_name": "relu", ...},

The # activations section comment should remain on its own line to preserve the visual grouping:

    {"op_name": "isnan", ...},
    {"op_name": "isinf", ...},
    # activations
    {"op_name": "relu", ...},

Fix this →

2. Hardcoded torch.float32 for the inf constant — backends/mlx/ops.py:468

inf_slot = emit_lifted_constant(P, float('inf'), torch.float32)

The inf constant is always created as float32, even when the input tensor x may be float16 or bfloat16. This works because EqualNode likely handles type promotion internally, but it's worth verifying this is intentional and doesn't cause issues for half-precision inputs. The test does cover float16 and bfloat16, so if CI passes this is likely fine — but a comment noting the intentional choice or using the input's dtype would make this clearer.

3. Minor: _inf_input_fn overwrites some +inf with -inf — backends/mlx/test/test_ops.py:4029-4032

mask_pos = torch.rand(shape) > 0.8   # ~20% set to +inf
mask_neg = torch.rand(shape) > 0.9   # ~10% set to -inf (may overwrite +inf)
x[mask_pos] = float('inf')
x[mask_neg] = float('-inf')

Since mask_neg is applied after mask_pos, some positions marked +inf will be overwritten to -inf. This isn't a bug per se (both infinities should be detected), but the effective distribution is slightly different from what the thresholds suggest. The existing _nan_input_fn uses a single mask, which is cleaner. Consider using non-overlapping masks if the intent is to have distinct proportions, or just note that overlap is acceptable.

---

Looks Good

• The decomposition abs(x) == inf is mathematically correct and handles both +inf and -inf.
• Properly uses make_tmp_slot() for the intermediate abs result, matching codebase conventions.
• Correctly validates args/kwargs with require_args and require_kwargs.
• Good placement under the "Numerical checks" section alongside _isnan_handler.
• Test configuration with _SHAPES_3 and multiple float dtypes provides reasonable coverage.

---

Issue (1) is the only one I'd consider a must-fix before merge. The # activations comment being on its own line is part of the existing code organization pattern in that file, and the current diff breaks that grouping.
|

[metascroy]

Can we add nan cases as well?

[metascroy]

Looks great @Jah-yee! Left a small nit and running CI now.

[metascroy]

Can you run the lintrunner?

pip install -r requirements-lintrunner.txt
lintrunner -a

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/mlx

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[Jah-yee]

Thanks for the thorough review, @nil-is-all!

思考路径

收到 3 条反馈，逐一分析：

1. # activations 注释格式 — 确实违反代码风格。_UNARY_OP_TESTS 列表中 section comment 都是单独成行，isinf 被插入时意外合并了。这破坏了文件的视觉分组。

2. float32 硬编码 dtype — MLX 的 EqualNode 内部做类型提升，所以 float32 的 inf 能正确匹配 float16/bfloat16 的输入。但原文没说明意图，reviewer 提的对，加个注释消除歧义。

3. mask overlap — 这是我自己写 _inf_input_fn 时的疏忽，mask_neg 在 mask_pos 之后应用，会覆盖部分 +inf。但从功能角度不影响测试（两种 inf 都测到了），reviewer 也说 "not a bug per se"。不过非重叠 mask 更干净，改了。

修改方案

Fix 1: 格式化

# Before
{"op_name": "isinf", ...},    # activations
{"op_name": "relu", ...

# After
{"op_name": "isinf", ...},
# activations
{"op_name": "relu", ...

Fix 2: dtype 注释

# Create inf constant (float32; EqualNode handles type promotion to match input dtype)
inf_slot = emit_lifted_constant(P, float('inf'), torch.float32)

Fix 3: 非重叠 mask

# Before: mask_neg 覆盖部分 mask_pos
mask_pos = torch.rand(shape) > 0.8
mask_neg = torch.rand(shape) > 0.9  # may overwrite +inf

# After: 完全独立的 mask
mask_pos = torch.rand(shape) > 0.8
mask_neg = (~mask_pos) & (torch.rand(shape) > 0.9)

Commit 已推送: 2cdde9d7f

Please re-review!

[metascroy]

@Jah-yee it looks like lintrunner is still failing in CI, e.g.,

 Warning (UFMT) format
    Run `lintrunner -a` to apply this patch.

    You can run `lintrunner -a` to apply this patch.

    465  465 |     )
    466  466 | 
    467  467 |     # Create inf constant (float32; EqualNode handles type promotion to match input dtype)
    467      |-    inf_slot = emit_lifted_constant(P, float('inf'), torch.float32)
         468 |+    inf_slot = emit_lifted_constant(P, float("inf"), torch.float32)
    469  469 | 
    470  470 |     # Compare abs(x) == inf
    471  471 |     out = P.make_or_get_slot(n)

Make sure you run the linter locally and then resubmit the PR with the linted code.

[kirklandsign]

@Jah-yee 请改一下这2个 linter

>>> Lint for backends/mlx/ops.py:

  Warning (UFMT) format
    Run `lintrunner -a` to apply this patch.

    You can run `lintrunner -a` to apply this patch.

    465  465 |     )
    466  466 | 
    467  467 |     # Create inf constant (float32; EqualNode handles type promotion to match input dtype)
    467      |-    inf_slot = emit_lifted_constant(P, float('inf'), torch.float32)
         468 |+    inf_slot = emit_lifted_constant(P, float("inf"), torch.float32)
    469  469 | 
    470  470 |     # Compare abs(x) == inf
    471  471 |     out = P.make_or_get_slot(n)



>>> Lint for backends/mlx/test/test_ops.py:

  Warning (UFMT) format
    Run `lintrunner -a` to apply this patch.

    You can run `lintrunner -a` to apply this patch.

    4028  4028 |         # Insert ~20% +inf and ~10% -inf using non-overlapping masks
    4029  4029 |         mask_pos = torch.rand(shape) > 0.8  # ~20% -> +inf
    4030  4030 |         mask_neg = (~mask_pos) & (torch.rand(shape) > 0.9)  # ~10% of remaining -> -inf
    4030       |-        x[mask_pos] = float('inf')
    4031       |-        x[mask_neg] = float('-inf')
          4031 |+        x[mask_pos] = float("inf")
          4032 |+        x[mask_neg] = float("-inf")
    4033  4033 |         return (x,)
    4034  4034 | 
    4035  4035 |     return fn

[linux-foundation-easycla]

• ✅ login: Jah-yee / name: Jah-yee (2cdde9d)
• ❌ The email address for the commit (810bc80) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.
• ❌ - login: @falkecl / name: Jah-yee. The commit (f811554) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[Jah-yee]

Thanks for the suggestion on adding nan cases! Supporting aten.isnan in the MLX backend would require a separate decomposed op handler (similar to the isinf handler in this PR) plus corresponding tests - best as a follow-up PR. For this PR (isinf), the lint fixes have been applied. Happy to open a follow-up issue if you'd like to track the isnan enhancement.