Jax计算框架是Google用来取代Tensorflow的新一代计算框架，这个框架使用类似pytorch的技术，但是在pytorch技术之上加入了更加强大的技术，但是这也导致该框架使用起来要比pytorch难一些，但是该框架的计算性能又比较优秀，因此依旧具有较大的吸引力。
Jax框架的性能优势主要体现在单机多卡GPU设备的自动并行上，Jax和TensorFlow同样有着编译器的优势，Jax可以通过简单的代码编写实现多卡并行的方式，在数据并行和模型并行的基础之上又实现了数据切割并行的方式，可以说Jax框架在原生水平上支持目前所有的并行操作，并且在并行的代码编写上又比pytorch更加简单，但是由于并行的复杂性，Jax的自动并行往往有着不好掌握的特点，虽然在编写上代码更加精炼，但是逻辑难度加大了，不容易掌握。
可以说，Jax对于pytorch来说最大的优点就是对于多卡设备，或者说是大规模计算的原生支持上，比如在Jax官方的讲解文档之中都是默认对8个TPU进行操作的，为实现同样的操作（代码可运行），我们也需要使用8个卡的GPU，这个要求虽然对于企业来说都是极为容易实现的，但是对于个人用户这个要求却是极为困难的。如果没有多卡的运行环境，我们也就无法体验Jax框架的特性。
为了方便体验和debug Jax框架的多卡自动计算的特性，Jax框架给出了多卡仿真的设置参数：
官方文档：