Skip to content

Lower cummax op#8491

Merged
qihqi merged 5 commits into
pytorch:masterfrom
zyy-martin:yaoyangz/cummax
Jan 10, 2025
Merged

Lower cummax op#8491
qihqi merged 5 commits into
pytorch:masterfrom
zyy-martin:yaoyangz/cummax

Conversation

@zyy-martin
Copy link
Copy Markdown
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 ManfeiBai and bhavya01 December 20, 2024 04:01
@miladm miladm added the lowering ATen Operation lowering label Dec 20, 2024
@zyy-martin
Copy link
Copy Markdown
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
Copy Markdown
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
@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

Labels

lowering ATen Operation lowering

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Lower cummax

4 participants