JAX
2 篇文章
- 如何扩展你的模型（10）：用 JAX 给 TPU 编程中英对照第 10 篇：jax.jit / pmap / shard_map、GSPMD、pallas kernel 等并行 API 到底怎么用，配可运行的例子。左栏原文，右栏译文。
- 如何扩展你的模型（9）：如何给 TPU 代码做性能剖析中英对照第 9 篇：真实模型永远比理论复杂——JAX + XLA 栈、TensorBoard profiler、trace viewer，以及怎么定位真实的性能问题。左栏原文，右栏译文。
2 篇文章