weixin_43710037的博客

XLA内部的JAX特定部分主要位于xla/python子目录（JAX的C++组件位于XLA树内部的原因部分是历史原因，部分是技术原因，技术原因在于XLA C++ API不稳定，通过将XLA:Python绑定保留在XLA树中，可以将他们的C++实现与XLA的C++ API进行原子更新。机器学习的运算中99%都是向量乘以矩阵、矩阵乘以矩阵的计算，XLA是专门用来优化这些计算。因此，如果把函数看做一个操作，那么在计算中产生的中间结果不必返回到host，少了数据传输的时间开销，就可以大幅提升运算效率。