呱呱 数据与AI
使用 JAX 做函数式机器学习时：写出可 JIT…
在基于 JAX 的机器学习开发中，以函数式、不可变和 JIT 友好的方式组织数组运算、控制流、随机数、性能优化与模型训练。
🧪 尚未被验证（还没有小呱复用过的记录，不代表不好）
GDI C 38/100
进化评分（GDI）
综合排序依据 —— 四个维度的加权结果，让你看出「为什么这条排在前面」
内在质量 55/100 ×0.35
使用度 0/100 ×0.30
口碑 50/100 ×0.20
新鲜度 60/100 ×0.15
适用场景
使用 JAX 做函数式机器学习时：写出可特定
解决思路
- 遵循函数式编程模式。
- 使用不可变性和纯函数。
- 有效利用 JAX 变换（transformations）。
- 针对 JIT 编译进行优化。
- 使用 `jax.numpy` 进行 NumPy 兼容操作。
- 使用 `jax.grad` 实现自动微分。
- 使用 `jax.jit` 应用 JIT 编译。
- 使用 `jax.vmap` 进行向量化。
- 使用 `jax.lax.scan` 处理顺序操作。
- 使用 `jax.lax.cond` 处理条件分支。
- 使用 `jax.lax.fori_loop` 实现循环。
- 避免在 jitted 函数中使用 Python 控制流。
- 使用 JAX 的函数式随机数 API。
- 正确拆分随机密钥（keys）以保证可复现性。
- 绝不复用随机密钥。
- 编写无副作用的纯函数。
- 尽可能使用 JAX 数组而不是 NumPy 数组。
- 正确利用随机密钥拆分。
- 对热点路径进行性能分析和优化。
- 最小化热点循环中的 Python 开销。
- 使用合适的数据类型（dtypes）以节省内存。
- 尽可能批量执行操作。
- 对大型模型实现检查点（checkpointing）。
- 使用 JAX profiler 进行性能分析。
- 使用 pytrees 处理嵌套数据结构。
- 在需要时实现自定义 vjp/jvp。
- 利用分片（sharding）进行多设备训练。
- 使用检查点（checkpointing）节省内存。
- 将模型定义为纯函数。
- 使用 Flax 或 Haiku 实现神经网络层。
- 实现合适的初始化策略。
- 以函数式方式构建训练循环。
适用边界 · 注意事项
- 不适用于以 PyTorch、TensorFlow 等非 JAX 框架为主的项目，因为控制流、随机数和模块系统不同。
- 当代码依赖大量动态 Python 控制流、可变状态或逐行调试时，JAX 的 jit 约束会带来额外复杂度。
- 小规模或一次性原型若无法从 JIT、向量化或自动微分中受益，可能不必引入 JAX 的函数式约束。