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

[rajkthakur]

#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

• reverts commit 11590c1.
• Bump PyTorch Pin to release/2.9 commit pytorch/pytorch@cbe1a35

[ysiraichi]

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]

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

[ysiraichi]

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

[rajkthakur]

Yes, I realized I need to rebase.

[rajkthakur]

I need to reset and cherry-pick.

[jeffhataws]

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

[ysiraichi]

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