自动微分：

JAX提供了强大的自动微分功能，支持前向和反向模式的自动微分。它提供了如jax.grad、jax.hessian、jax.jacfwd和jax.jacrev等函数，用于计算数值函数的梯度、Hessian矩阵等。

这种自动微分功能使得研究人员能够轻松地处理机器学习中的优化问题，无需手动编写复杂的求导代码。

向量化与并行化：

JAX的vmap变换提供了自动矢量化算法，使得研究人员能够在大规模的数据上运行相同的函数，如计算整个批次的损失或每个样本的损失等。

此外，JAX还支持通过pmap转换实现大规模的数据并行，从而优雅地将单个处理器无法处理的大数据进行处理。

即时编译（JIT）：

JAX使用XLA（Accelerated Linear Algebra，加速线性代数）进行即时编译（JIT），能够在GPU和云TPU加速器上高效执行JAX程序。

JIT编译与JAX的API（与Numpy一致的数据函数）结合，为研发人员提供了便捷接入高性能计算的可能，无需特别的经验就能将计算运行在多个加速器上。

与Numpy的兼容性：

JAX的API基于Numpy构建，包含丰富的数值计算与科学计算函数，因此Python和Numpy的广泛使用使得JAX十分简洁、灵活、易于上手，学习成本也比较低。