🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139939

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job, 4 Pending

As of commit 393a3bd with merge base ada43ed ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@pytorchbot label "topic: not user facing"

Any review on this?

Thanks for your implementation! I have a question regarding the shape check for xs:

assert x.shape == shape, "All xs tensors must have the same shape"

Why does it require the tensors to have exactly the same shape? In the JAX implementation, only the first dimension is required to match:

num_elems = int(elems_flat[0].shape[axis])
if not all(int(elem.shape[axis]) == num_elems for elem in elems_flat[1:]):

Hi @WeihanLikk

Thank you for looking into this. I was slow in working on this over the holidays but will pick up steam again. You are right, I don't think that this is necessarily required. Just the scanned dimension needs to be identical for all xs. I will take a look

@garrett361 I've implemented a first version of the backward approach we discussed offline. The algorithm per se works, but there are still some things to sort out. In particular the lifted argument and partial gradient support.

EDIT: Of course there is still further room to cleanup the code and to adjust the documentation.

lifted argument and partial gradient support

What does lifted argument mean?

Partial gradient = when only some inputs require_grad?

What does lifted argument mean?

In some cases, variables and other properties from the combine_fn are lifted as additional inputs. For example, these could be external variables, or symbolic shapes of tensors as well. In the following case

H = torch.rand(2, device=device)
def fct_freevars1(x: torch.Tensor, y: torch.Tensor):
    return x * H + y * 2

H would become a lifted variable. I know how to handle those, but I think @ydwu4 is currently working on simplifying the autograd architecture for higher order operators, such as associative scan, to simplify this handling.

Partial gradient = when only some inputs require_grad?

Yes, that is correct. Same applies here as well. I know how to handle it, but I wanted to wait for the autograd rework.

Currently using associative scan for a big research project related to linear attention. (or perhaps I should say, logarithmic attention 😉)

Is there an expected timeline for autograd support to be available? Really excited about this PR!!

@huydhn I wouldn't know of this failure, but it could be of course

@pytorchbot merge -f 'Incorrect revert, this change is not related to the trunk failure'

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Merge failed

Reason: PR #139939 has not been reviewed yet

@pytorchbot merge -f 'Incorrect revert, this change is not related to the trunk failure'

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team