Hard dependency on `jax[tpu]` blocks or degrades installation on all non-TPU platforms

Tunix declares `jax[tpu]` as a hard dependency in `pyproject.toml`. The `[tpu]` extra pulls in `libtpu`, which only ships Linux wheels — so the package is uninstallable on macOS, and on other non-TPU platforms it installs an unneeded TPU runtime with real side effects.

## Who is affected

This affects every non-TPU platform, not just macOS:

- **macOS**: `libtpu` ships no darwin wheels, so `pip install google-tunix` fails outright.
- **Linux + GPU**: the forced `libtpu` install is not just dead weight (~100MB+). Its presence makes JAX attempt TPU backend initialization (noisy errors, sometimes requiring a manual `JAX_PLATFORMS=cuda` override), and the `jax[tpu]` version constraint can silently re-resolve the user's `jax` to a version that no longer matches their `jax-cuda12-plugin`, breaking the GPU backend entirely.
- **Linux + CPU** (including your own `cpu-tests.yml` CI): works, but installs an unused TPU runtime on every run.

In other words, the hard `[tpu]` pin costs something on every platform except TPU hosts — where it saves only the typing of an extra suffix.

## Why CPU installs matter

The common workflow for agentic RL development is to iterate and verify small-scale on an ordinary dev machine before moving to TPU. Tunix's own CI relies on exactly this: `cpu-tests.yml` runs `pytest tests/rl/agentic/` (and most of `tests/rl/`) on `ubuntu-latest` with no accelerator. On Linux the `[tpu]` extra happens to resolve (installing an unused `libtpu`), which is why this isn't visible in CI — but on macOS it makes the package uninstallable.

**Verified:** installing tunix with `--no-deps` plus plain CPU `jax` on macOS (arm64), the full agentic stack (`ToolAgent`, `ToolEnvironment`, `TrajectoryCollectEngine`) imports and collects complete rollout episodes correctly.

## Proposal

Happy to send a PR for whichever direction you prefer:

1. Depend on plain `jax>=0.6.0,!=0.7.2,<0.11.1` and add a `tpu` extra (`pip install "google-tunix[tpu]"`), matching the convention of flax/optax; or
2. Keep `jax[tpu]` but gate it with an environment marker (`platform_system == "Linux"`) and add plain `jax` for other platforms.

Option 1 is cleaner but changes behavior for existing TPU users installing the bare package; option 2 is fully backward-compatible.