cannot import name 'pytree' from 'jaxlib'

[HouJun19]

Hi I'm trying to run this code on my ubuntu18.04,but It turns out

Traceback (most recent call last):
  File "ilqr_jax_MPC.py", line 1, in <module>
    from jax import jit, jacfwd, jacrev, hessian, lax
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/__init__.py", line 19, in <module>
    from jax.api import *
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/api.py", line 37, in <module>
    from . import core
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 28, in <module>
    from . import dtypes
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/dtypes.py", line 31, in <module>
    from .lib import xla_client
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/lib/__init__.py", line 55, in <module>
    from jaxlib import pytree
ImportError: cannot import name 'pytree' from 'jaxlib' (/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jaxlib/__init__.py)

Actually I have run this code successfully on jaxlib 0.1.47,but It turns out ti be very slow to complete one loop which needs more than 2 hours, and it has a warning that :Can not found GPU, fall back to CPU.
So I want to use GPU to accelerate the code, while I have tried many versions but all failed. For my GPU is NVIDIA 3060ti, some old version of jaxlib and CUDA are not support for the gpu, when I try the jaxlib 0.1.52 I have the same erro that " cannot import name 'pytree' from 'jaxlib' ". while using newer jaxlib and jax it also turns out some other erros like

pygame 2.1.2 (SDL 2.0.16, Python 3.8.13)
Hello from the pygame community. https://www.pygame.org/contribute.html
Traceback (most recent call last):
File "ilqr_jax_MPC.py", line 164, in
jac_l, hes_l, jac_l_final, hes_l_final, jac_f = derivative_init()
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/src/api.py", line 527, in cache_miss
out_flat = xla.xla_call(
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 1937, in bind
return call_bind(self, fun, *args, **params)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 1953, in call_bind
outs = top_trace.process_call(primitive, fun, tracers, params)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 687, in process_call
return primitive.impl(f, *tracers, **params)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py", line 208, in _xla_call_impl
compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/linear_util.py", line 295, in memoized_fun
ans = call(fun, *args)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py", line 257, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py", line 302, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 2188, in trace_to_jaxpr_final2
jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 2139, in trace_to_subjaxpr_dynamic2
out_tracers = map(trace.full_raise, ans)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/util.py", line 47, in safe_map
return list(map(f, *args))
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 415, in full_raise
return self.pure(val)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1761, in new_const
aval = raise_to_shaped(get_aval(c), weak_type=dtypes.is_weakly_typed(c))
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 1163, in get_aval
return concrete_aval(x)
File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 1155, in concrete_aval
raise TypeError(f"Value {repr(x)} with type {type(x)} is not a valid JAX "
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Value <CompiledFunction of <function jacfwd..jacfun at 0x7fe3d73ac670>> with type <class 'jaxlib.xla_extension.CompiledFunction'> is not a valid JAX type

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

I really don't know how to solve them. Sorry if I have some misunderstanding because I just start using python and still a lack of experience.
Looking forward to your reply.

[jakevdp]

We've seen this error before, and it is generally due to one of two things:

1. You're using a newer version of the jax package with an older version of the jaxlib package (Newer jax attempts to import jaxlib.pytree, which didn't exist in older jaxlib versions)
2. You're using an incorrectly-built version of the jaxlib package (jaxlib.pytree is a compiled extension, and will not import if the package is not compiled properly)

You report that you're using very old versions of jax and jaxlib (v0.1.X); I'd suggest updating to more recent releases if possible. The current releases are jax v0.3.15 and jaxlib 0.3.15.