Enhance Operator Robustness and Expand Supported Ops

[srmurali002]

Enhance Operator Robustness and Expand Supported Ops

Summary

This PR introduces a robust version fallback mechanism for operator handlers and expands support for several key ONNX operators, including a new implementation of HardSigmoid. These changes improve the library's future-proofing and compatibility with modern ONNX models (e.g., RTMPose).

Key Changes

1. Robust Version Fallback Mechanism (handler.py)

• Refactored Handler.handle to implement a highest-supported-version fallback.
• If an exact match for a requested opset version isn't implemented, the handler now automatically selects the highest available version $\le$ the requested version.
• Benefit: Prevents NotImplementedError when running models with newer opsets (e.g., opset 22/25) that haven't changed specific operator semantics since their last implemented version.

2. New Operator: HardSigmoid

• Added full implementation of the HardSigmoid operator using JAX.
• Registered version 6 handler in hardsigmoid.py.

3. Operator Fixes and Version Updates

• Clip:
  • Added version_11 support.
  • Fixed argument handling to correctly support min and max as positional inputs (consistent with newer ONNX specs).
• ReduceSum:
  • Added version_11 support.
• Squeeze:
  • Added version_11 support.
  • Fixed axes attribute handling to ensure consistency with ONNX expectations for opset 11+.

4. Standards and Quality

• Strong Typing: Added comprehensive type hints to all new/modified functions and classes.
• Linting: All modified files pass pylint with a 10.00/10 score, adhering strictly to the Google Python Style Guide.
• Circular Imports: Cleaned up internal imports to maintain a stable module initialization.

Verification Results

Unit Tests

• Added test_hardsigmoid_ patterns to tests/onnx_ops_test.py.
• Verified that all patched operators (HardSigmoid, Clip, ReduceSum, Squeeze) pass the standard ONNX backend tests on TPU.

Checklist

• Code follows the Google Python Style Guide.
• All modified files pass pylint.
• New operator implementation includes unit tests.
• Verified compatibility with JAX backend on TPU.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.