MPAX 是一个基于 JAX 的数学编程库，旨在使深度学习中的数学编程变得轻松。它支持硬件加速（CPU、GPU 和 TPU），并且是可微的，可以直接集成到神经网络的反向传播中。此外，MPAX 还支持批处理和分布式求解，显著提高了大规模问题的计算效率。
