JAX 0.2.17与0.4.1版本内存占用差异过大问题咨询

针对StackOverflow上的Python性能优化问题，使用其采纳回答中的jit_bar函数进行测试，代码如下：

import numpy as np
import jax
import jax.numpy as jnp

T = np.random.rand(5000, 566, 3)
@jax.jit
def jit_bar(Y):
   u, v = jnp.triu_indices(Y.shape[0], 1)
   return jnp.sqrt((3 * (Y[u] - Y[v]) ** 2).mean(axis=(-1, -2)))
msd = jit_bar(T)

当传入维度为(10000, 566, 3)的数组时，Python 3.6环境下内存占用稳定在1.5GB左右，但在Python 3.8/3.9/3.10/3.11环境中，内存占用骤升至50GB以上。

经测试确认，该问题仅与JAX版本相关：

• 正常运行环境：Python 3.6、JAX 0.2.17、jaxlib 0.1.68、numpy 1.19.2
• 内存异常环境：Python 3.11、JAX 0.4.1、jaxlib 0.4.1、numpy 1.24.1

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