torch_tpu: update pin to 2b237048ad2d20bb083bc1eebb7b73b1a4d3daa0

.github/actions/build-torch-tpu/action.yml

@@ -16,7 +16,7 @@ runs:
           --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
           --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
           --pre \
-          'jax==0.10.2' 'jaxlib==0.10.2' 'libtpu==0.0.43' 'tpu-info==0.7.1' 'jaxtyping' 'frozendict' 'immutabledict'
+          'jax==0.11.0' 'jaxlib==0.11.0' 'libtpu==0.0.44.1' 'tpu-info==0.14.2' 'jaxtyping' 'frozendict' 'immutabledict'
         # Install Bazel
         if ! command -v bazel &> /dev/null; then
           sudo curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64 -o /usr/local/bin/bazel

.github/actions/install-pytorch-for-torchtpu/action.yml

@@ -13,4 +13,4 @@ runs:
         source .venv/bin/activate
         # The latest PyTorch nightly known to build against the pinned torch_tpu
         # commit (.github/ci_commit_pins/torch_tpu.txt). Bump the two together.
-        uv pip install -U --pre "torch==2.14.0.dev20260707" --index-url https://download.pytorch.org/whl/nightly/cpu
+        uv pip install -U --pre "torch==2.14.0.dev20260726" --index-url https://download.pytorch.org/whl/nightly/cpu

.github/ci_commit_pins/torch_tpu.txt

@@ -1 +1 @@
-058f3a3e4099f472728e3d52676b53ce411d02ca
+9d8945e064e0145c154cd510866182d58363d9ba

.github/workflows/test.yml

@@ -205,7 +205,7 @@ jobs:
         if: matrix.alias == 'pallas-interpret'
         run: |
           source .venv/bin/activate
-          uv pip install 'jax==0.10.1' 'jaxlib==0.10.1' 'absl-py'
+          uv pip install 'jax==0.11.0' 'jaxlib==0.11.0' 'absl-py'
 
       - name: Install CUTLASS CuTe DSL
         if: matrix.backend == 'cute'