高级指南# 本节包含关于多核计算、自动微分和自定义操作等高级主题的示例和教程。 并行计算 分布式数组和自动并行化 显式分片(亦称“类型内分片”) 使用 shard_map 进行手动并行化 设备本地数组布局控制 JAX 内存和主机卸载 优化器状态卸载 多控制器 JAX 简介(即多进程/多主机 JAX) 分布式数据加载 自动微分 自动微分秘籍 自定义导数规则 使用 jax.checkpoint(亦称 jax.remat)控制自动微分的保存值 深入探讨 JAX 中的广义卷积 XLA 编译器标志