AttributeError: module 'jax' has no attribute 'linear_util'

[shriva1]

Can anyone please help with this error?

AttributeError Traceback (most recent call last)
Cell In[3], line 3
1 import sys
2 import jax
----> 3 import jax_md
5 print("Python Version:", sys.version)
6 print("JAX Version:", jax.version)

File c:\Users\aishw\anaconda3\envs\lgnn\Lib\site-packages\jax_md_init_.py:16
1 # Copyright 2019 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
(...)
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
15 from jax_md import space
---> 16 from jax_md import energy
17 from jax_md import minimize
18 from jax_md import simulate

File c:\Users\aishw\anaconda3\envs\lgnn\Lib\site-packages\jax_md\energy.py:28
26 from jax.tree_util import tree_map
27 from jax import vmap
---> 28 import haiku as hk
...
52 warnings.warn(message, DeprecationWarning, stacklevel=2)
53 return fn
---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}")