Skip to content

Revert XLA backward incompatible breaking change - #9668 - Fixes Issue #9685#9686

Merged
qihqi merged 2 commits intopytorch:r2.9from
rajkthakur:rajkthakur/revert-openxla-breaking-change
Oct 22, 2025
Merged

Revert XLA backward incompatible breaking change - #9668 - Fixes Issue #9685#9686
qihqi merged 2 commits intopytorch:r2.9from
rajkthakur:rajkthakur/revert-openxla-breaking-change

Conversation

@rajkthakur
Copy link
Copy Markdown

@rajkthakur rajkthakur commented Oct 21, 2025 •
edited
Loading

#9668 upgraded OpenXLA: openxla/xla@92f7b59 to openxla/xla@9a9aa0e
which introduced a change 5bba026 related to instructions ID that is backward incompatible and causing neuron to crash - Issue #9685

Changes include

@pytorch-bot pytorch-bot Bot added the ci-no-td label Oct 21, 2025
@jeffhataws jeffhataws requested review from bhavya01, qihqi, ysiraichi and zhanyong-wan and removed request for bhavya01 October 21, 2025 17:43
ysiraichi
ysiraichi approved these changes Oct 21, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@ysiraichi ysiraichi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall, looks good to me.
I have a few minor comments:

  1. Could you be more specific describing the changes in this PR? e.g.:
    • Specify which commit you are reverting
    • Say that, even though you are reverting that commit, you are bumping the PyTorch pin
  2. Let's wait, at least, 1 day for @bhavya01 input.

@jeffhataws
Copy link
Copy Markdown
Collaborator

jeffhataws commented Oct 21, 2025

@rajkthakur should this go to r2.9 branch instead? We still want to move to newer openxla for top-of-tree I think.

@rajkthakur rajkthakur changed the base branch from master to r2.9 October 21, 2025 20:51
@ysiraichi
Copy link
Copy Markdown
Collaborator

ysiraichi commented Oct 21, 2025 •
edited
Loading

@rajkthakur Could you rebase this on top of r2.9?

@rajkthakur
Copy link
Copy Markdown
Author

rajkthakur commented Oct 21, 2025

Yes, I realized I need to rebase.

@rajkthakur
Copy link
Copy Markdown
Author

rajkthakur commented Oct 21, 2025

I need to reset and cherry-pick.

@rajkthakur rajkthakur force-pushed the rajkthakur/revert-openxla-breaking-change branch from 5a44e19 to 37a0305 Compare October 21, 2025 23:34
@jeffhataws
Copy link
Copy Markdown
Collaborator

jeffhataws commented Oct 22, 2025

@ysiraichi @bhavya01 @qihqi is the torchax failure a concern?

@ysiraichi
Copy link
Copy Markdown
Collaborator

ysiraichi commented Oct 22, 2025

I'm not sure. But, since it's being pushed to the release branch, I think it's important to get everything right.
Try this patch (also, call it out in the PR description):

diff --git a/torchax/dev-requirements.txt b/torchax/dev-requirements.txt
index 2da02ae85..e6315be57 100644
--- a/torchax/dev-requirements.txt
+++ b/torchax/dev-requirements.txt
@@ -2,4 +2,5 @@
 torch==2.8.0 ; sys_platform == 'darwin'  # macOS
 torch==2.8.0+cpu; sys_platform != 'darwin' # Non-macOS (CPU-only), like on TPU
 yapf==0.40.2  # N.B.: keep in sync with `infra/ansible/config/pip.yaml`, `.github/workflows/lintercheck.yml`
+jax==0.7.2  # N.B.: torchax breaks on newer JAX versions that would be pulled from `flax` dependencies
 flax==0.10.6

@qihqi qihqi merged commit abc7854 into pytorch:r2.9 Oct 22, 2025
21 of 24 checks passed
@jeffhataws jeffhataws mentioned this pull request Oct 31, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@qihqi qihqi qihqi approved these changes

@ysiraichi ysiraichi ysiraichi approved these changes

@bhavya01 bhavya01 Awaiting requested review from bhavya01

@zhanyong-wan zhanyong-wan Awaiting requested review from zhanyong-wan

Assignees

No one assigned

Labels

ci-no-td

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@rajkthakur @jeffhataws @ysiraichi @qihqi