[torchlib] Consolidate all overloads and prevent new ones from being created

[justinchuby]

This PR implements #2580 by combining all overloads in torchlib and remove the ability to register new ones. It is done in a BC compatible fashion and should work with released versions of PyTorch.

From now on all logic for a single aten OpOverload should be implemented by a single torchlib function to ensure 1-to-1 mapping.

[Copilot]

Pull Request Overview

This PR consolidates overloaded functions in the torch_lib by removing the concept of private functions and preventing new overloads from being created for the same operation name. The changes focus on simplifying the registration system and merging boolean indexing operations with their regular counterparts.

• Removes the private parameter and functionality from the registration system
• Consolidates boolean and regular index operations into unified functions
• Adds validation to prevent duplicate overload registrations

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxscript/function_libs/torch_lib/registration.py	Removes private function support and adds overload duplication prevention
onnxscript/function_libs/torch_lib/ops/core.py	Consolidates index and index_put functions to handle both boolean and integer indexing
tests/function_libs/torch_lib/ops_test_data.py	Removes separate boolean index test entries and duplicates
onnxscript/function_libs/torch_lib/ops/nn.py	Removes outdated comment about private functions
onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py	Updates to exclude removed private functions from iteration

[titaiwangms]

CI is failing

[codecov]

Codecov Report

❌ Patch coverage is 82.35294% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.13%. Comparing base (726be2b) to head (e6c4551).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/registration.py	60.00%	4 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2621      +/-   ##
==========================================
- Coverage   70.14%   70.13%   -0.02%     
==========================================
  Files         226      226              
  Lines       27369    27379      +10     
  Branches     2775     2779       +4     
==========================================
+ Hits        19199    19201       +2     
- Misses       7218     7223       +5     
- Partials      952      955       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

In general, to fix “explicit returns mixed with implicit returns” in a function that is annotated to return a non-None value, ensure that every control-flow path either explicitly returns a value of the annotated type or raises an exception. You prevent implicit fall-through by adding an explicit return at the end of the function or at the end of any branch where fall-through is possible.

For _aten_index_bool, the best fix without changing functionality is to make sure the else: branch (for non-all-scalar boolean indices) also ends with an explicit return of a TensorType. Inside this branch, the function presumably computes or delegates to _aten_index_onnx in a way analogous to the scalar-boolean case. We should reuse that result instead of inventing new behavior. Concretely, we can have the else: branch finish by returning the value from whatever computation is already there. If the omitted code builds some processed_indices and currently calls _aten_index_onnx(self, processed_indices, index_ranks) without returning it, we should change that call into return _aten_index_onnx(...). If, instead, the omitted code only prepares data and then falls through, we can add a final return _aten_index_onnx(self, indices, index_ranks) at the end of the function so that any path in the else: branch still returns a TensorType. This keeps semantics aligned with the aten_index function (which always delegates to _aten_index_onnx).

Given we must not assume unshown code structure, the minimally invasive and safe change within the shown region is to add an explicit return _aten_index_onnx(self, indices, index_ranks) as the last statement of _aten_index_bool. That guarantees that if execution reaches the bottom of the else: branch (or any other path that doesn’t otherwise return/raise), the function will still return a TensorType rather than None. No new imports or helper methods are needed; _aten_index_onnx is already used within this module.