Skip to content

Conversation

justinchuby
Copy link
Collaborator

Previously sizes coming in as SymInt[] are first concatenated as INT64 then used. This created inefficiencies where we could not process any static dims from the size list and had to treat the whole shape as dynamic. In aten_expand, this meant we needed to add Abs on the shape.

This change updates the functions that take SymInt[] such that they are no longer turned into INT64 first. I updated aten_expand to process constant -1 values so an Abs is not required. I also added a helper merge_dims to create constants for consecutive constant dims first before concatinating.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby justinchuby added the module: torchlib Related to the torch/aten function lib in development label Aug 27, 2025
@justinchuby justinchuby requested a review from Copilot August 27, 2025 16:42
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR improves the handling of SymInt[] parameters in torch_lib operations by avoiding unnecessary conversion to INT64 tensors and enabling better optimization of static dimensions.

  • Updates functions to accept Sequence[INT64] instead of IntType for shape parameters
  • Introduces a merge_dims helper function to optimize consecutive constant dimensions
  • Removes unnecessary Abs operations in aten_expand by handling constant -1 values directly

Reviewed Changes

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

File Description
onnxscript/function_libs/torch_lib/ops/core.py Updates multiple torch operations to use new shape handling approach and removes unnecessary type casting
onnxscript/function_libs/torch_lib/ops/common.py Adds new merge_dims helper function to optimize consecutive constant dimensions

Copy link

codecov bot commented Aug 27, 2025

Codecov Report

❌ Patch coverage is 93.93939% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.98%. Comparing base (a925acc) to head (4334c37).
⚠️ Report is 2 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
onnxscript/function_libs/torch_lib/ops/core.py 91.66% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2522      +/-   ##
==========================================
+ Coverage   69.93%   69.98%   +0.04%     
==========================================
  Files         216      216              
  Lines       26052    26058       +6     
  Branches     2616     2618       +2     
==========================================
+ Hits        18219    18236      +17     
+ Misses       6931     6920      -11     
  Partials      902      902              

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

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby justinchuby added this to the 0.4.1 milestone Aug 27, 2025
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
s
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby justinchuby enabled auto-merge (squash) August 28, 2025 19:24
@justinchuby justinchuby merged commit fc792e4 into main Sep 4, 2025
32 checks passed
@justinchuby justinchuby deleted the justinchu/symint branch September 4, 2025 17:16
@github-project-automation github-project-automation bot moved this from Todo to Done in ONNX Script Review Board Sep 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: torchlib Related to the torch/aten function lib in development
Projects
Development

Successfully merging this pull request may close these issues.

4 participants