Add torchax maximum version. (#9706)

This fixes the `torchax` CI failures (e.g. [#9705 action][1]) caused by
the [`flax` dependency][2], which pulls `jax` 0.8.0 (or newer) for some
Python versions.

**Key Changes:**

- Set upper bound `<0.8.0` in `pyproject.toml` optional dependency
section
- Install `torchax` before the test dependencies

[1]:
https://github.com/pytorch/xla/actions/runs/19229289257/job/54963835669?pr=9705
[2]:
https://github.com/pytorch/xla/blob/611a5cc2675133e6b159a2be8b07b65c44656f29/torchax/dev-requirements.txt#L5

Author: ysiraichi

.github/workflows/torchax.yml

@@ -46,8 +46,8 @@ jobs:
         shell: bash
         working-directory: torchax
         run: |
-          pip install -r test-requirements.txt
           pip install -e .[cpu]
+          pip install -r test-requirements.txt
       - name: Run tests
         if: needs.check_code_changes.outputs.has_code_changes == 'true'
         working-directory: torchax

torchax/pyproject.toml

@@ -40,11 +40,11 @@ classifiers = [
 path = "torchax/__init__.py"
 
 [project.optional-dependencies]
-cpu = ["jax[cpu]>=0.6.2", "jax[cpu]"]
+cpu = ["jax[cpu]>=0.6.2, <0.8.0"]
 # Add libtpu index `-f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html`
-tpu = ["jax[cpu]>=0.6.2", "jax[tpu]"]
-cuda = ["jax[cpu]>=0.6.2", "jax[cuda12]"]
-odml = ["jax[cpu]>=0.6.2", "jax[cpu]"]
+tpu = ["jax[cpu,tpu]>=0.6.2, <0.8.0"]
+cuda = ["jax[cpu,cuda12]>=0.6.2, <0.8.0"]
+odml = ["jax[cpu]>=0.6.2, <0.8.0"]
 
 [tool.hatch.build.targets.wheel]
 packages = ["torchax"]