Skip to content

Add support for PyTorch non-finite predicates - #2834

Open
binaydhakal wants to merge 3 commits into
apple:mainfrom
binaydhakal:add-torch-nonfinite-predicates
Open

Add support for PyTorch non-finite predicates#2834
binaydhakal wants to merge 3 commits into
apple:mainfrom
binaydhakal:add-torch-nonfinite-predicates

Conversation

@binaydhakal

Copy link
Copy Markdown
Contributor

Summary

Add PyTorch frontend support for:

  • torch.isinf
  • torch.isfinite
  • torch.isposinf
  • torch.isneginf

The lowering compares floating-point inputs with dtype-matched positive and negative infinity constants. This avoids arithmetic-based detection that can misclassify very large finite FP32 values when FLOAT16 compute precision is enabled.

The implementation also preserves PyTorch semantics for FP16, integer, and boolean tensors. isinf and isfinite support complex values by evaluating their real and imaginary components; PyTorch itself rejects the signed-infinity predicates for complex inputs.

Both TorchScript and Torch Export are covered.

Testing

  • 48 new native conversion/prediction cases covering all four predicates, TorchScript and Torch Export, ML Program FP16, neural-network FP32, FP16/FP32/int32/bool inputs, extreme finite values, NaN, signed infinity, and complex inputs
  • 40 existing TestNanToNum regression cases
  • 104 existing complex and logical-operator regression cases
  • Separate eight-configuration native Core ML runtime matrix: FP32/FP16 inputs x FP32/FP16 compute precision x CPU-only/all compute units

All tests passed on macOS.

Comment thread coremltools/converters/mil/frontend/torch/ops.py Outdated
Signed-off-by: Binaya Dhakal <binaydhakal35@gmail.com>
@TobyRoseman

Copy link
Copy Markdown
Collaborator

@TobyRoseman

Copy link
Copy Markdown
Collaborator

There are CI failures. Please take a look. It seems some of the model aren't valid for EXECUTORCH, in which case you should skip those tests for that frontend. There are lots of examples of that in the same test file.

Signed-off-by: Binaya Dhakal <binaydhakal35@gmail.com>
@binaydhakal

binaydhakal commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @TobyRoseman, the remaining PR-related failures were the four complex-input cases under the ExecuTorch frontend. aten.complex.default is not part of ExecuTorch's Core ATen opset, so I skipped TestNonFinitePredicates.test_complex for ExecuTorch using the same guard already used by TestComplex.

Local validation:

  • ExecuTorch complex predicate cases: 4 skipped
  • TorchScript TestNonFinitePredicates: 24 passed
  • Torch Export TestNonFinitePredicates: 24 passed

@TobyRoseman

Copy link
Copy Markdown
Collaborator

@binaydhakal

Copy link
Copy Markdown
Contributor Author

It passed this time @TobyRoseman . Thank you so much for your help and guidance.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants