📘 JAX · TF 融合指南
30章
🎒 风格 · 活力色系
第 1 章
JAX与TensorFlow生态融合概述
⭐
为什么需要融合 · 核心概念对比 · 典型应用场景
第 2 章
JAX基础快速回顾
grad / jit / vmap / pmap · 函数式编程 · 不可变数组与随机数
第 3 章
TensorFlow基础快速回顾
TF 2.x核心API · Eager Execution · Keras · SavedModel
第 4 章
数据桥接
JAX NumPy ↔ tf.Tensor · 性能开销与最佳实践
第 5 章
模型桥接
jax2tf导出 · tf.function包装 · 控制流与副作用
第 6 章
训练桥接
TF优化器+JAX梯度 · JAX循环+TF数据 · 混合训练
第 7 章
序列化与部署
SavedModel · TF Serving · TFLite部署JAX模型
第 8 章
jax2tf深入
转换规则 · 支持类型 · 常见错误与调试
第 9 章
TensorFlow Probability与JAX
TFP JAX后端 · 概率分布 · 贝叶斯推断与MCMC
第 10 章
Flax与Keras集成
Flax→Keras层 · Keras循环+Flax · 混合回调
第 11 章
Haiku与Keras集成
Haiku→Keras模型 · 复用预训练权重 · 实践案例
第 12 章
Optax与TF优化器
互操作性 · TF循环+Optax · 梯度裁剪与调度
第 13 章
数据加载与预处理
tf.data→JAX · 文本/图像预处理 · 性能对比
第 14 章
分布式训练融合
pmap vs MirroredStrategy · 混合使用 · 梯度同步
第 15 章
XLA编译器的角色
共享XLA后端 · 优化融合模型 · 调试编译错误
第 16 章
自定义操作与扩展
JAX调用TF自定义op · TF调用JAX函数 · 性能
第 17 章
调试与日志融合
TensorBoard监控JAX · Summary写入 · 统一日志
第 18 章
模型量化与优化
TFLite量化JAX · 量化感知训练 · 边缘部署
第 19 章
强化学习框架融合
TF-Agents+JAX策略 · JAX调用TF环境 · 实践
第 20 章
自然语言处理融合
HuggingFace TF+JAX · 加载TF权重 · 混合推理
第 21 章
计算机视觉融合
Keras预训练视觉模型 · JAX图像pipeline · 实践
第 22 章
时间序列与信号处理
复用TF信号函数 · 时序数据集+JAX · 性能对比
第 23 章
生成模型融合
加载TF GAN权重 · Keras层构建JAX生成器 · 混合训练
第 24 章
图神经网络融合
TF-GNN + JAX · 数据格式转换 · 性能优化
第 25 章
科学计算融合
TF数值积分/微分 · 优化器求解JAX方程 · 实践
第 26 章
生产环境部署
TFX Pipeline部署JAX · 版本管理 · A/B测试
第 27 章
性能基准测试
JAX vs TF vs 融合 · 内存分析 · 计算图优化
第 28 章
社区生态与工具
Orbax / TensorStore · 最佳实践 · 学习资源
第 29 章
未来趋势与展望
融合方向 · JAX在TF生态定位 · 对AI框架影响
第 30 章
综合实战项目
端到端JAX-TF融合 · 数据加载/训练/部署 · 全知识点