Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Lower cummax op #8491

Merged
merged 5 commits into from
Jan 10, 2025
Merged

Lower cummax op #8491

merged 5 commits into from
Jan 10, 2025

Conversation

zyy-martin
Copy link
Contributor

@zyy-martin zyy-martin commented Dec 13, 2024

Implement cummax op. Fixes #8230

Cummax returns a tuple of (values, indices). This PR uses the ReduceWindowWithGeneralPadding XLA op (padding + window reducer for cumulative values) to implement it similar to cumsum/cumprod. The only difference is that the indices will be passed to the ReduceWindowWithGeneralPadding op along with the values, and a new reducer is implemented to look at both the value and index of input and return a tuple of (max_value, max_index).

@zyy-martin zyy-martin marked this pull request as ready for review December 18, 2024 19:21
@zyy-martin zyy-martin changed the title implement cummax op Lower cummax op Dec 18, 2024
@miladm miladm requested review from bhavya01 and ManfeiBai December 20, 2024 04:01
@zyy-martin
Copy link
Contributor Author

I exclude the cummax from the python op test because I believe the scalar value generated from the sample input would fail, as will existing ops like cumprod and cumsum. The following script will fail for example:

import torch
import torch_xla
import torch_xla.core.xla_model as xm

a = torch.tensor(7.358110427856445, device=xm.xla_device())
print(f'[xla]: {torch.cumsum(a, 0)}')

@zyy-martin
Copy link
Contributor Author

Thanks for the review! Can someone with write access merge the PR? Thanks!

@qihqi qihqi merged commit c394d1b into pytorch:master Jan 10, 2025
12 checks passed
@zyy-martin zyy-martin deleted the yaoyangz/cummax branch January 10, 2025 21:45
@zyy-martin zyy-martin mentioned this pull request Jan 14, 2025
qihqi pushed a commit that referenced this pull request Jan 16, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Lower cummax
4 participants