环境准备与工具链搭建：JAX、PyTorch、Flax、Haiku、Trax 等库的安装与版本兼容性检查

说实话，做模型互转这件事，最让人头疼的往往不是模型本身，而是环境。

我见过太多人，代码写得漂漂亮亮，结果一跑就报错。一看，版本冲突。嗯，这坑我踩过不止一次。所以这一章，咱们先把地基打牢。

1. 核心库概览：你要装什么？

先捋清楚我们要装哪些东西。说白了，就是两套生态：

• PyTorch 生态：PyTorch 本身，加上 torchvision（图像）、torchtext（文本）这些辅助库。
• JAX 生态：JAX 是底层，Flax、Haiku、Trax 是上层框架。

我个人习惯把这两套装在不同的虚拟环境里。为什么？你想想看，JAX 对 CUDA 版本特别敏感，PyTorch 也是。混在一起，很容易出现「装好了 JAX，PyTorch 跑不了」的尴尬局面。

2. 版本兼容性：最容易翻车的地方

我曾经在一个客户现场，花了两小时排查一个 bug。最后发现是 JAX 0.3.x 和 Flax 0.6.x 不兼容。嗯，从那以后，我每次装环境都会先查版本对照表。

下面这张表，是我自己整理的，比较常用且稳定的组合：

JAX 版本	Flax 版本	Haiku 版本	Trax 版本	PyTorch 版本
0.4.x	0.7.x	0.0.9+	1.4.x	2.0.x
0.3.x	0.6.x	0.0.8	1.3.x	1.13.x
0.2.x	0.5.x	0.0.7	1.2.x	1.10.x

3. 安装命令：一行一行来

我建议你按顺序装。别一次性 pip install 一长串，出了问题不好定位。

3.1 安装 PyTorch

# CPU 版本
pip install torch torchvision torchaudio

# GPU 版本（以 CUDA 11.8 为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

装完之后，跑一下验证：

import torch
print(torch.__version__)
print(torch.cuda.is_available())  # 应该返回 True

3.2 安装 JAX

JAX 的安装稍微特殊一点。官方推荐用 jaxlib 分开装。

# CPU 版本
pip install jax jaxlib

# GPU 版本（CUDA 11.8）
pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

3.3 安装上层框架

# Flax
pip install flax

# Haiku
pip install haiku

# Trax
pip install trax

装完之后，验证一下：

import jax
import flax
import haiku as hk
import trax

print(jax.__version__)
print(flax.__version__)
print(hk.__version__)
print(trax.__version__)

4. 知识体系：一张图看懂

下面这张 SVG 图，是我画的环境与工具链关系图。你可以把它当作「地图」，后面每章都会用到。

5. 避坑指南：我踩过的那些坑

• JAX 的 GPU 版本装不上？ 我曾经在 Ubuntu 20.04 上折腾了一整天。最后发现是 libcudnn 版本不对。解决方案：用 conda install -c conda-forge jaxlib 代替 pip。
• Flax 和 Haiku 能混用吗？ 理论上可以，但我不建议。它们的设计哲学不同。Flax 更像 PyTorch 的 nn.Module，Haiku 更像函数式编程。混用容易精神分裂。
• Trax 装完报错？ 检查一下 tensorflow 是否装了。Trax 依赖 TensorFlow 的数据集 API。嗯，这个坑我帮三个同事解决过。

6. 验证脚本：一键检查

我写了个小脚本，每次装完环境跑一遍，心里踏实：

import sys
import subprocess

packages = ['torch', 'jax', 'flax', 'haiku', 'trax']

for pkg in packages:
    try:
        mod = __import__(pkg)
        version = getattr(mod, '__version__', 'unknown')
        print(f'✅ {pkg}: {version}')
    except ImportError:
        print(f'❌ {pkg}: 未安装')
    except Exception as e:
        print(f'⚠️ {pkg}: 导入出错 - {e}')

跑完之后，如果全是绿色勾勾，恭喜你，环境搭好了。

好了，环境准备好了。接下来我们就可以开始真正的模型互转了。

---