tensor.roll() #3281

crutcher · 2025-06-10T22:06:09Z

tensor.roll()

Checklist

Confirmed that cargo run-checks command has been executed.
Made sure the book is up to date with changes in this PR.

Related Issues/PRs

See pytorch roll:
https://docs.pytorch.org/docs/stable/generated/torch.roll.html

Changes

Added burn::tensor::indexing
Added ReversableIndex (for usize / isize dims)
Added roll(), roll_dim()

Testing

Gen tests.

crutcher · 2025-06-11T18:22:40Z

Probably needs some more tests; but I'd like a review of the ReflectableIndex code and if changes needed to be made there

codecov · 2025-06-11T18:34:52Z

Codecov Report

Attention: Patch coverage is 89.55696% with 33 lines in your changes missing coverage. Please review.

Project coverage is 82.73%. Comparing base (81985bd) to head (6c1583a).

Files with missing lines	Patch %	Lines
crates/burn-tensor/src/tensor/indexing/mod.rs	83.21%	23 Missing ⚠️
crates/burn-tensor/src/tensor/api/base.rs	90.47%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3281      +/-   ##
==========================================
+ Coverage   82.66%   82.73%   +0.06%     
==========================================
  Files         995      997       +2     
  Lines      127626   127933     +307     
==========================================
+ Hits       105498   105839     +341     
+ Misses      22128    22094      -34

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

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

laggui

Some thoughts about the indexing stuff:

ReflectableIndex seems very much in line with the IndexConversion trait (previously added to extend the slicing/indexing support). It should probably be unified, and perhaps AsIndex would be more idiomatic for a name.

Now, I am a bit conflicted about using the same trait for element indexing and dimension/axis selection. While technically we can consider them both forms of indexing into different structural levels of a tensor, it might be a little confusing. And for a public API, perhaps it is better to separate them. The same canonicalizing function could still be used under the hood to handle negative values based on the size of a dimension (for AsIndex) vs the rank of a tensor (for AsAxis/AsDim). Or we could have a parent trait common to both. But for the public API, it might make more sense to differentiate them. Let me know what you think.

Finally, the number of places by which elements are shifted in roll and roll_dim are semantically different. shift is just a type that support negative values to wrap around, so not sure that the interface should have shift: ReflectableIndex. It's only being used to convert to isize and then wrap the negative values based on the size.

crutcher · 2025-06-13T18:51:14Z

copied from Discord thread

I would like to lift the trait and utility code into burn::indexing; as we'll probably have a lot of related support machinery, including ravel/unravel support with bounds checking, as that is always part of tensor binding. Particularly for things like binding external datatypes (such as Images) into/out-of TensorData.

I do want to open a disagreement about having separate semantic types; and clarify a strategy here.

I am on board, in general, with the "new type" pattern to do param / unit separation, particularly in situations where we don't have keyword params; but I don't think that applies here. We're not discussing explicit casting into the appropriate semantic type 'dimension, index, shape' at the call location; in fact most of the work is to permit usize and isize to be used interchangeably without additional call-site syntax.

There is some protection if we make the utitily machinery require the newtypes, so you can't call something like wrap_index(size, index) and swap the params; but there's also a higher expectation of "know what I'm doing" when calling the indexing utilities explicitly; they will generally only be used by lower level implementors. And your comments aren't proposing that change anyway.

There's some additional (potential) utility in labeling the panic errors better; but those could be addressed by passing in labels to the utility mechanisms.

My preference would be to lift IndexConversion into burn::indexing, and not have seperate types for dim, index, size, and extent (what shift is here); just because, with the current auto-convert plans, the other approach puts a lot of weight on implementors without having any resulting api change for callers

crutcher · 2025-06-13T19:56:49Z

Per discussion with @laggui :

Renamed IndexConversion to AsIndex
Added more trait impls for the basic integer types
moved to burn::indexing
also, reordered roll fns to : (dim, shift) param order.
added tests

laggui

Looks pretty good to me! I agree with the changes as discussed 👍

My only comments are mostly related to convention / practices.

Also, we should probably make the canonicalize_* functions return a result instead of panicking internally (esp. for external usage). And in roll we can unwrap.

crates/burn-tensor/src/lib.rs

crates/burn-tensor/src/tensor/api/base.rs

laggui · 2025-06-16T14:03:04Z

crates/burn-tensor/src/tensor/api/base.rs

+    pub fn roll_dim<I1, I2>(self, dim: I2, shift: I1) -> Self
+    where
+        I1: AsIndex,
+        I2: AsIndex,


Even if both generics are bound to AsIndex, I would use better / more descriptive names than I1 and I2. Since D is already used for the const rank generic, maybe Dim and Shift.

Also, I thought you wanted to change the order to tensor.roll_dim(shift, dim) (discord)

... my memory of the outcome of that discussion doesn't match the thread; I did extra work to make this change, because it used to be the other way ...

I'll change it back.

Could also be my memory 😅 but I think dim is typically after the input arguments to specify on which dim/axis the operation is applied. And that's how most of these ops are currently defined (though there might be some inconsistencies as you pointed out last time).

crates/burn-tensor/src/tensor/api/base.rs

crates/burn-tensor/src/tensor/indexing/mod.rs

laggui · 2025-06-16T14:18:44Z

crates/burn-tensor/src/tests/ops/roll.rs

+    fn test_roll_empty() {
+        let device = Default::default();
+        let input = TestTensorInt::<2>::zeros([12, 0], &device);
+
+        // Rolling an empty tensor should return the same empty tensor
+        input
+            .clone()
+            .roll(&[0, 1], &[1, 2])
+            .to_data()
+            .assert_eq(&input.to_data(), false);
+    }


I tested this use case locally for CUDA, which doesn't have the same restrictions on 0 size resources as wgpu. We have a separate issue when returning the data for a buffer of size 0.

Probably need to handle this special case for cubecl into_data method when bytes.is_empty(). Can take care of this before merging.

crates/burn-tensor/src/tensor/indexing/mod.rs

laggui

Hmm not sure why the macos CI hangs.. might actually be related to the empty test.

Will have to investigate before we merge. Otherwise, just some minor comments.

laggui · 2025-06-17T14:08:06Z

crates/burn-tensor/src/tensor/indexing/mod.rs

+/// ## Arguments
+///
+/// * `idx` - The dimension index to canonicalize.
+/// * `rank` - The size of the dimension.


rank != size of the dimension

laggui · 2025-06-17T14:12:10Z

crates/burn-tensor/src/tensor/api/base.rs

+                size,
+            );
+            assert!(
+                dim < size,


Didn't catch this in the previous round, is this check correct? size here is the size of the dim.. why are we checking that the dimension specified is smaller than its size?

- Move IndexConversion to burn::indexing - Reorder to (dim, shift) - improve error messages.

…tests

crutcher changed the title ~~[WIP] tensor.roll()~~ tensor.roll() Jun 11, 2025

laggui reviewed Jun 12, 2025

View reviewed changes

crutcher force-pushed the crutcher/roll branch from 2b1401b to 7b13e92 Compare June 13, 2025 20:04

laggui reviewed Jun 16, 2025

View reviewed changes

crutcher force-pushed the crutcher/roll branch from 636648a to 1d2eade Compare June 16, 2025 19:15

crutcher requested a review from laggui June 16, 2025 19:27

laggui reviewed Jun 17, 2025

View reviewed changes

crutcher requested a review from laggui June 17, 2025 19:30

crutcher added 16 commits June 17, 2025 14:23

WIP tensor.roll()

ff14802

docs

e7467dd

doc

96ebf1b

use slice_dim

9df0386

might_panic for empty tensors

4fcadc6

std gate

66163c1

cleanup params

b0ad999

docs

5cf6846

- Unify around IndexConversion

4df020b

- Move IndexConversion to burn::indexing - Reorder to (dim, shift) - improve error messages.

Rename IndexConversion to AsIndex (per discussion with @laggui); add …

e2dbf8e

…tests

Fix logic bug in roll

1e1789d

Simplify unchecked logic; naming

79e2d4e

Review changes

8a6838f

final review changes

44b8002

ordering

d3fef9a

review

6c1583a

crutcher force-pushed the crutcher/roll branch from 9af8e30 to 6c1583a Compare June 17, 2025 21:34

tensor.roll() #3281

Are you sure you want to change the base?

tensor.roll() #3281

Conversation

crutcher commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklist

Related Issues/PRs

Changes

Testing

Uh oh!

crutcher commented Jun 11, 2025

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

laggui left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crutcher commented Jun 13, 2025

copied from Discord thread

Uh oh!

crutcher commented Jun 13, 2025

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

crutcher commented Jun 10, 2025 •

edited

Loading

codecov bot commented Jun 11, 2025 •

edited

Loading

laggui left a comment •

edited

Loading