2026年1月4日
JAX性能优化实战：7个变换让TPU/GPU吃满算力
摘要： JAX跑得快的技巧其实很简单：通过组合变换让XLA能看到大块连续的计算，比如说批处理、融合、分片，让每一步在单设备或多设备同步时都像一个干净的kernel。 我们今天就来总结7个能够提高运行速度的JAX变换组合 https://avoid.overfit.cn/post/84e4e28e3ca847 阅读全文
posted @ 2026-01-04 20:44 deephub 阅读(3) 评论(0) 推荐(0)