feat(atenlib): add op(chunk) #511
Conversation
Codecov Report
@@ Coverage Diff @@
## main #511 +/- ##
==========================================
+ Coverage 72.72% 72.79% +0.06%
==========================================
Files 109 109
Lines 10689 10703 +14
Branches 1102 1104 +2
==========================================
+ Hits 7774 7791 +17
+ Misses 2611 2607 -4
- Partials 304 305 +1
... and 1 file with indirect coverage changes Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. |
@gramalingam pls help to review. The original implementation doesn't return valid output numbers and based on opset17(split). |
| num_per_chunk = op.Div(dim_size, chunks) | ||
| if op.Mod(dim_size, chunks) > 0: # type: ignore[operator] | ||
| # cannot be divisible, then num_in_chunk + 1 | ||
| num_per_chunk = op.Add(num_per_chunk, 1) |
There was a problem hiding this comment.
I wonder if using Where will be more efficient or not
There was a problem hiding this comment.
Something like op.Where(op.Mod(dim_size, chunks) > 0, 1, 0) + num_per_chunk
There was a problem hiding this comment.
Even better: op.Cast(op.Mod(dim_size, chunks) > 0, to=INT64.value) + num_per_chunk
There was a problem hiding this comment.
The take away for me was If tends to be slow and should be replaced with other things when possible.
|
Also cc @BowenBao for the logic |
justinchuby
added
the
module: torchlib
|
We can merge if there's no movement |
No description provided.