这是一个结合了即时编译(JIT)、自动求导(Autograd)、线性代数编译器(XLA) 的数值计算 Python 库，用法和 NumPy 差不多。但 JAX 更快、更省内存，支持自动微分、自动向量化、并行计算等功能。