关于JAX处理百万级独立小型非线性方程组的技术咨询

我目前采用MPI将百万个独立小型非线性方程组分配到200个CPU处理，每个CPU依次求解5000个系统，过程包含微分运算与线性代数求解。现针对JAX相关能力咨询以下问题：

• JAX求解单个非线性方程组效率不错，但处理百万级任务时，推荐的并行执行方式是什么？我疑惑GPU/TPU是否主要提升单任务性能，而非多并发任务？
• 当非线性方程组维度低于100时，JAX与NumPy的性能对比如何？JAX的优势是否主要体现在高维场景？

更新说明 #

我的代码为多模块PDE求解程序，核心流程如下：

n  = 1000000
m  = 100
y  =  jnp.zeros((n, m))

for i in jnp.arange(1,n):
    y = y.at[i, :].set(jnp.linalg.solve(A, b))

其中A和b由其他函数计算，各迭代的A、b输入来自大型数据集且相互独立。我了解Numba的prange可并行化该循环，想知道JAX是否有类似功能。此外我测试了MPI4JAX，将循环分配到10个CPU核心实现了10倍加速，MPI类并行随核心数增加线性提速，但工作站CPU核心限100，每个核心仍需顺序处理10000次迭代，期望利用GPU实现更快处理。

内容的提问来源于stack exchange，提问作者funpy