Update jax requirement from <=0.4.37,>=0.3.2 to >=0.3.2,<=0.4.38

Update jax requirement from <=0.4.37,>=0.3.2 to >=0.3.2,<=0.4.38 #41