machine-learning 这份 Skill 聚焦 JAX 与函数式模型开发。它把数组计算、自动微分、编译和向量化组织成明确的函数与数据流。
随机数也要显式管理键，拆出不同用途的键，避免无意重复同一序列。可复现还需核输入、版本和硬件条件，并非记住一个种子就保证全流程一致。
主文建议用 JAX 控制流处理编译函数中的条件与循环；静态 Python 控制并非全都禁止，关键是值是否在追踪时确定。编译也有首次成本和形状变化影响，不能只凭加 JIT 就说训练更快。
适合先理清状态与计算边界，再测热点。本文未训练模型或运行代码，材料不是覆盖数据验证、评估和部署的全部机器学习流程；多设备切分与检查点也需结合实际内存和通信成本。