Merge pull request #2 from abhijeetgangan/Fix

Fix config import in test_util

jax_md/test_util.py

@@ -25,7 +25,7 @@
 
 import jax.numpy as jnp
 import numpy as onp
-from jax.config import config
+from jax import config
 
 import jax
 from jax import dtypes as _dtypes