高级指南# 本节包含有关更高级主题的示例和教程,例如多核计算、自动微分和自定义操作。 并行计算 分布式数组和自动并行化 使用 shard_map 进行手动并行化 多主机和多进程环境 分布式数据加载 自动微分 自动微分食谱 自定义导数规则 使用 jax.checkpoint (又名 jax.remat) 控制自动微分保存的值 深入探讨 JAX 中的广义卷积 XLA 编译器标志列表