JAX性能优化实战：7个变换让TPU/GPU吃满算力

JAX跑得快的技巧其实很简单：通过组合变换让XLA能看到大块连续的计算，比如说批处理、融合、分片，让每一步在单设备或多设备同步时都像一个干净的kernel。

我们今天就来总结7个能够提高运行速度的JAX变换组合

 

https://avoid.overfit.cn/post/84e4e28e3ca8473488a0e9248d1ec51b