Add Blackwell / sm_100 (B200) support: CUDA 12.8 + jax 0.4.38

[maxscheurer]

Summary

The pinned torch==2.7.0+cu126 and jax/jaxlib==0.4.29 cannot run on NVIDIA
Blackwell GPUs (compute capability 10.0 / sm_100, e.g. B200). This PR bumps the
environment to CUDA 12.8 + jax 0.4.38 so the stack builds and runs on Blackwell,
while remaining compatible with existing Ampere/Hopper/Ada setups.

Closes #54 and #36.

Why the current pins fail on Blackwell

Component	Current pin	Problem on sm_100
PyTorch	2.7.0+cu126	cu126 wheels ship kernels only up to sm_90 → CUDA error: no kernel image is available for execution on the device.
JAX	0.4.29	Predates Blackwell; jaxlib 0.4.29 cannot target sm_100.

Changes

• env/build_uv_env.sh
  • torch → 2.7.0+cu128 (cu128 wheels are a superset of cu126's arch list,
    so no regression on sm_70–sm_90 GPUs) and PyG wheel index → cu128.
  • Pin torchvision==0.22.0+cu128 / torchaudio==2.7.0+cu128 to the versions
    matching torch 2.7.0. Leaving them unpinned lets uv resolve a newer torchaudio
    (e.g. 2.11.0) whose ABI is incompatible with torch 2.7.0
    (undefined symbol: torch_library_impl), which breaks import transformers
    / ESMFold in the evaluate stage.
  • jax[cuda12]/jaxlib → 0.4.38 (same minor series as the previously validated
    0.4.29 → minimal numerical drift for the AF2 oracle, but Blackwell-capable),
    plus dm-haiku==0.0.13. CUDA 12.8 libraries come from the PyPI nvidia-*
    wheels via --index-strategy unsafe-best-match.
• pyproject.toml: [tool.uv] extra-index cu126 → cu128.
• community_models/colabdesign/__init__.py: restore the top-level
  jax.tree_map/tree_flatten/tree_unflatten/tree_leaves aliases (removed in
  jax ≥ 0.4.31) from jax.tree_util. Guarded by hasattr, so it is a no-op
  on older jax where the aliases still exist — safe for anyone not upgrading jax.

Compatibility

• cu128 covers sm_70–sm_120, so existing Ampere (A100), Hopper (H100), and
  Ada (L40S) users are unaffected by the torch bump.
• The colabdesign shim only adds attributes that are missing, so it changes
  nothing on the previous jax 0.4.29.
• Requires a CUDA 12.8-capable driver (R525+ via minor-version compatibility).

Testing

Built and verified end-to-end on both B200 (sm_100) and L40S (sm_89):

• Clean env/build_uv_env.sh --clean build; torch + jax import and run GPU ops.
• colabdesign / AF2 modules import (shim active on jax 0.4.38).
• A full complexa design run (generate → filter → evaluate → analyze) with the
  AF2 reward/refold oracle, producing valid (non-NaN) metrics.