Implement aten::repeat_interleave operators #2464

Copilot · 2025-07-24T16:43:49Z

This PR implements the three variants of aten::repeat_interleave that were previously stubbed out with NotImplementedError:

repeat_interleave.Tensor(Tensor repeats, *, SymInt? output_size=None) -> Tensor
repeat_interleave.self_Tensor(Tensor self, Tensor repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor
repeat_interleave.self_int(Tensor self, SymInt repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor

Implementation Details

The implementation uses a combination of ONNX operations to achieve the repeat_interleave behavior:

Core Algorithm: Uses CumSum + ArgMax + Less operations to implement a "searchsorted" equivalent in ONNX, efficiently mapping output positions to input element indices
Key ONNX Operations: CumSum, ArgMax, Range, Gather, Tile, Reshape, Less, Cast

Examples

# repeat_interleave.Tensor - creates index pattern
torch.repeat_interleave(torch.tensor([1, 2, 3]))
# Output: tensor([0, 1, 1, 2, 2, 2])

# repeat_interleave.self_Tensor - repeats elements with variable counts  
torch.repeat_interleave(torch.tensor([10, 20, 30]), torch.tensor([2, 1, 3]))
# Output: tensor([10, 10, 20, 30, 30, 30])

# repeat_interleave.self_int - repeats elements with fixed count
torch.repeat_interleave(torch.tensor([1, 2, 3]), 2)
# Output: tensor([1, 1, 2, 2, 3, 3])

# All variants support dim parameter for multi-dimensional tensors
torch.repeat_interleave(torch.tensor([[1, 2], [3, 4]]), 2, dim=0)
# Output: tensor([[1, 2], [1, 2], [3, 4], [3, 4]])

Testing

Added comprehensive test entries to ops_test_data.py for all three variants
Verified functionality with manual tests covering edge cases including:
- Basic 1D tensor operations
- Multi-dimensional tensors with dim parameter
- Zero repeat counts
- Single element tensors
All implementations pass linter checks (RUFF, MYPY, PYLINT, EDITORCONFIG-CHECKER)

Fixes #2462.

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

Co-authored-by: justinchuby <[email protected]>

codecov · 2025-07-24T17:28:40Z

❌ 23 Tests Failed:

Tests completed	Failed	Passed	Skipped
16458	23	16435	3853

View the top 3 failed test(s) by shortest run time

::onnxscript.tools.training_helper

Stack Traces | 0s run time

ImportError while importing test module '.../onnxscript/tools/training_helper.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

::onnxscript.tools.transformers_models.llama_test

Stack Traces | 0s run time

ImportError while importing test module '.../tools/transformers_models/llama_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../tools/transformers_models/llama_test.py:12: in <module>
    import onnxscript.tools.training_helper
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

::onnxscript.tools.transformers_models.mistral_test

Stack Traces | 0s run time

ImportError while importing test module '.../tools/transformers_models/mistral_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../tools/transformers_models/mistral_test.py:14: in <module>
    import onnxscript.tools.training_helper
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

tests/function_libs/torch_lib/ops_test_data.py

xadupre · 2025-07-24T18:27:52Z

onnxscript/function_libs/torch_lib/ops/core.py

+        # Repeat each index 'repeats' times
+        indices_reshaped = op.Unsqueeze(original_indices, [1])  # Shape: [dim_size, 1]
+        repeat_pattern = op.Constant(value_ints=[1, repeats])
+        repeated_indices = op.Tile(


The Tile should apply on self, not repeated the indices to gather the right element later. It is slow.

Optimized the implementation to use Tile directly on the tensor instead of creating indices and using Gather. This eliminates the intermediate index operations and improves performance as suggested. Changes made in commit a80b64b.

xadupre · 2025-07-24T18:28:45Z

onnxscript/function_libs/torch_lib/ops/core.py

+        # We can use Tile with appropriate reshaping
+        indices_reshaped = op.Unsqueeze(original_indices, [1])  # Shape: [num_elements, 1]
+        repeat_pattern = op.Constant(value_ints=[1, repeats])
+        repeated_indices = op.Tile(


Same comment here as below.

Updated the implementation to use a cleaner code structure similar to the optimized self_int version. Improved comments and variable naming for better consistency and readability in commit d291ae3.

xadupre · 2025-07-24T18:31:05Z

onnxscript/function_libs/torch_lib/ops/core.py

+) -> TensorType:
+    """repeat_interleave.self_Tensor(Tensor self, Tensor repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor"""
+
+    if dim is None:


No Tile here but the code should be very similar.

Optimized the aten_repeat_interleave_self_tensor function with a cleaner code structure and approach similar to the self_int version. While we can't use Tile directly with variable repeats, the implementation now follows the same pattern and style for consistency in commit d291ae3.

It is worse.

You're absolutely right. I reverted the aten_repeat_interleave_self_tensor function back to the original clean implementation from commit 6d62024. The verbose comments and variable name changes I added were unnecessary and made the code worse without providing any actual improvements. The original implementation was already clean and correct.

Fixed in commit a22cbfb - now aten_repeat_interleave_self_tensor uses the original clean approach while aten_repeat_interleave_self_int keeps the proper Tile optimization.

You can't use Argmax or ScatterND.

Removed ArgMax and ScatterND operations as requested. Replaced ArgMax with ReduceSum-based approach for the searchsorted algorithm and ScatterND with Concat-based pattern building. The implementation now uses only allowed ONNX operations while maintaining the same functionality. Fixed in commit 5409040.

…stead of indices Co-authored-by: xadupre <[email protected]>

…d improved comments Co-authored-by: xadupre <[email protected]>

onnxscript/function_libs/torch_lib/ops/core.py

…timization Co-authored-by: xadupre <[email protected]>

…ntations Co-authored-by: xadupre <[email protected]>

onnxscript/function_libs/torch_lib/ops/core.py

+
+        # Get the rank and build tile pattern
+        rank = op.Size(original_shape)
+        ones_before = op.ConstantOfShape(


onnxscript/function_libs/torch_lib/ops/core.py

+            op.Constant(value_ints=[1]),
+        )
+        repeat_val = op.Constant(value_ints=[repeats])
+        ones_after = op.ConstantOfShape(


Initial plan

e075cfa

Copilot AI assigned Copilot and justinchuby Jul 24, 2025

github-project-automation bot added this to ONNX Script Review Board Jul 24, 2025

github-project-automation bot moved this to Todo in ONNX Script Review Board Jul 24, 2025

Copilot started work on behalf of justinchuby July 24, 2025 16:43 View session

Copilot AI and others added 2 commits July 24, 2025 17:06

Implement repeat_interleave functions and add test entries

3ee33f5

Co-authored-by: justinchuby <[email protected]>

Fix code formatting and pass all linters

9f7d6dc

Co-authored-by: justinchuby <[email protected]>

Copilot AI changed the title ~~[WIP] Implement aten::repeat_interleave~~ Implement aten::repeat_interleave operators Jul 24, 2025

Copilot finished work on behalf of justinchuby July 24, 2025 17:19

Copilot AI requested a review from justinchuby July 24, 2025 17:19

justinchuby reviewed Jul 24, 2025

View reviewed changes

tests/function_libs/torch_lib/ops_test_data.py Outdated Show resolved Hide resolved

Update tests/function_libs/torch_lib/ops_test_data.py

d4d9805

justinchuby reviewed Jul 24, 2025

View reviewed changes

tests/function_libs/torch_lib/ops_test_data.py Outdated Show resolved Hide resolved

Update tests/function_libs/torch_lib/ops_test_data.py

6d62024

Copilot started work on behalf of justinchuby July 24, 2025 18:00 View session

Copilot finished work on behalf of justinchuby July 24, 2025 18:06

Copilot AI requested a review from justinchuby July 24, 2025 18:06

xadupre reviewed Jul 24, 2025

View reviewed changes

Copilot started work on behalf of xadupre July 24, 2025 18:27 View session

xadupre reviewed Jul 24, 2025

View reviewed changes

Optimize repeat_interleave.self_int to use Tile directly on tensor in…

3e4bf1b

…stead of indices Co-authored-by: xadupre <[email protected]>

Copilot AI requested a review from xadupre July 24, 2025 18:43

Copilot finished work on behalf of xadupre July 24, 2025 18:43

Copilot started work on behalf of xadupre July 24, 2025 18:43 View session

Optimize repeat_interleave.self_Tensor with cleaner code structure an…

d291ae3

…d improved comments Co-authored-by: xadupre <[email protected]>

Copilot finished work on behalf of xadupre July 24, 2025 18:59

github-advanced-security bot found potential problems Jul 24, 2025

View reviewed changes

onnxscript/function_libs/torch_lib/ops/core.py Fixed Show fixed Hide fixed

Copilot started work on behalf of xadupre July 25, 2025 05:40 View session

Revert self_tensor to original clean implementation, keep self_int op…

a22cbfb

…timization Co-authored-by: xadupre <[email protected]>

Copilot finished work on behalf of xadupre July 25, 2025 05:51

Copilot started work on behalf of xadupre July 25, 2025 09:08 View session

Remove ArgMax and ScatterND operations from repeat_interleave impleme…

5409040

…ntations Co-authored-by: xadupre <[email protected]>

Copilot finished work on behalf of xadupre July 25, 2025 09:24

github-advanced-security bot found potential problems Jul 26, 2025

View reviewed changes

xadupre mentioned this pull request Aug 6, 2025

Implements repeat_interleave #2477

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Implement aten::repeat_interleave operators #2464

Implement aten::repeat_interleave operators #2464

Uh oh!

Copilot AI commented Jul 24, 2025 •

edited

Loading

Uh oh!

codecov bot commented Jul 24, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

xadupre Jul 24, 2025

Uh oh!

Copilot AI Jul 24, 2025

Uh oh!

xadupre Jul 24, 2025

Uh oh!

Copilot AI Jul 24, 2025

Uh oh!

xadupre Jul 24, 2025

Uh oh!

Copilot AI Jul 24, 2025

Uh oh!

xadupre Jul 25, 2025

Uh oh!

Copilot AI Jul 25, 2025

Uh oh!

xadupre Jul 25, 2025

Uh oh!

Copilot AI Jul 25, 2025

Uh oh!

Uh oh!

Check failure

Check failure

Uh oh!

Implement aten::repeat_interleave operators #2464

Are you sure you want to change the base?

Implement aten::repeat_interleave operators #2464

Uh oh!

Conversation

Copilot AI commented Jul 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Implementation Details

Examples

Testing

Uh oh!

codecov bot commented Jul 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 23 Tests Failed:

Uh oh!

Uh oh!

Uh oh!

xadupre Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

xadupre Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

xadupre Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 24, 2025

Choose a reason for hiding this comment

Uh oh!

xadupre Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

xadupre Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Check failure

Check failure

Uh oh!

Copilot AI commented Jul 24, 2025 •

edited

Loading

codecov bot commented Jul 24, 2025 •

edited

Loading