使用 JAX 做函数式机器学习时：写出可 JIT…

在基于 JAX 的机器学习开发中，以函数式、不可变和 JIT 友好的方式组织数组运算、控制流、随机数、性能优化与模型训练。

适用场景

解决思路

1. 遵循函数式编程模式。
2. 使用不可变性和纯函数。
3. 有效利用 JAX 变换（transformations）。
4. 针对 JIT 编译进行优化。
5. 使用 `jax.numpy` 进行 NumPy 兼容操作。
6. 使用 `jax.grad` 实现自动微分。
7. 使用 `jax.jit` 应用 JIT 编译。
8. 使用 `jax.vmap` 进行向量化。
9. 使用 `jax.lax.scan` 处理顺序操作。
10. 使用 `jax.lax.cond` 处理条件分支。
11. 使用 `jax.lax.fori_loop` 实现循环。
12. 避免在 jitted 函数中使用 Python 控制流。
13. 使用 JAX 的函数式随机数 API。
14. 正确拆分随机密钥（keys）以保证可复现性。
15. 绝不复用随机密钥。
16. 编写无副作用的纯函数。
17. 尽可能使用 JAX 数组而不是 NumPy 数组。
18. 正确利用随机密钥拆分。
19. 对热点路径进行性能分析和优化。
20. 最小化热点循环中的 Python 开销。
21. 使用合适的数据类型（dtypes）以节省内存。
22. 尽可能批量执行操作。
23. 对大型模型实现检查点（checkpointing）。
24. 使用 JAX profiler 进行性能分析。
25. 使用 pytrees 处理嵌套数据结构。
26. 在需要时实现自定义 vjp/jvp。
27. 利用分片（sharding）进行多设备训练。
28. 使用检查点（checkpointing）节省内存。
29. 将模型定义为纯函数。
30. 使用 Flax 或 Haiku 实现神经网络层。
31. 实现合适的初始化策略。
32. 以函数式方式构建训练循环。

适用边界 · 注意事项

• 不适用于以 PyTorch、TensorFlow 等非 JAX 框架为主的项目，因为控制流、随机数和模块系统不同。
• 当代码依赖大量动态 Python 控制流、可变状态或逐行调试时，JAX 的 jit 约束会带来额外复杂度。
• 小规模或一次性原型若无法从 JIT、向量化或自动微分中受益，可能不必引入 JAX 的函数式约束。