Make JAX optional for tinker with skyrl-train backend

[tyler-griggs]

Summary

This PR refactors skyrl-tx to allow running the tinker API server with the skyrl-train backend without requiring JAX dependencies.

Stacks on top of #999.

Changes

1. Made JAX Optional (pyproject.toml)

• Removed jax>=0.8,<1.0, flax>=0.12.2, optax>=0.2.5 from base dependencies
• Moved all three to the [jax] extra only
• Base install no longer requires JAX

2. Separated Loss Type Definitions (types.py)

• Added LOSS_TYPES mapping to types.py (no JAX imports)
• This is just a string-to-index dict for validation and backend dispatch
• Added comment to keep it in sync with loss_fns.py

3. Lazy Load Backends (engine.py)

• Removed eager imports of JaxBackend and SkyRLTrainBackend at module level
• Created get_backend_classes() function that imports backends on-demand
• Backends now only loaded when --backend=<name> is specified
• Changed isinstance(self.backend, JaxBackend) to self.config.backend == "jax" to avoid import

4. Updated Import (engine.py)

• Changed from tx.tinker.loss_fns import LOSS_TYPES to from tx.tinker.types import LOSS_TYPES
• No more JAX dependency in engine.py

Clean Separation Verified

No JAX imports in:

• tx/tinker/types.py
• tx/tinker/engine.py
• tx/tinker/api.py
• tx/tinker/db_models.py
• tx/tinker/config.py
• tx/tinker/backends/skyrl_train.py
• tx/tinker/backends/backend.py
• tx/tinker/backends/utils.py

JAX only in backend-specific files:

• tx/tinker/backends/jax.py
• tx/tinker/loss_fns.py

Installation Patterns

# Tinker API with skyrl-train backend (NO JAX!)
pip install skyrl-tx[tinker,skyrl_train]
python -m tx.tinker.api --backend=skyrl_train --base_model=Qwen/Qwen3-0.6B

# Tinker API with JAX backend
pip install skyrl-tx[tinker,jax]
python -m tx.tinker.api --backend=jax --base_model=Qwen/Qwen3-0.6B

# Both backends available
pip install skyrl-tx[tinker,jax,skyrl_train]

Test plan

• Verify tinker API starts with skyrl-train backend without JAX installed
• Verify JAX backend still works with JAX installed
• Verify LOSS_TYPES mappings are in sync between types.py and loss_fns.py
• Run existing tinker tests with both backends

🤖 Generated with Claude Code

[tyler-griggs]

Closed in favor of #1004 which is based on main instead of tyler/tinker-sampling-main, to avoid including the sampling changes.