<h3 dir="auto">Description</h3>
<p dir="auto">Hi I'm trying to run this <a href="https://github.com/YukunXia/Carla_iLQR_MPC">code</a> on my ubuntu18.04,but It turns out</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;ilqr_jax_MPC.py&quot;, line 1, in &lt;module&gt;
    from jax import jit, jacfwd, jacrev, hessian, lax
  File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/__init__.py&quot;, line 19, in &lt;module&gt;
    from jax.api import *
  File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/api.py&quot;, line 37, in &lt;module&gt;
    from . import core
  File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 28, in &lt;module&gt;
    from . import dtypes
  File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/dtypes.py&quot;, line 31, in &lt;module&gt;
    from .lib import xla_client
  File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/lib/__init__.py&quot;, line 55, in &lt;module&gt;
    from jaxlib import pytree
ImportError: cannot import name 'pytree' from 'jaxlib' (/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jaxlib/__init__.py)"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "ilqr_jax_MPC.py", line 1, in &lt;module&gt;
    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 &lt;module&gt;
    from jax.api import *
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/api.py", line 37, in &lt;module&gt;
    from . import core
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py", line 28, in &lt;module&gt;
    from . import dtypes
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/dtypes.py", line 31, in &lt;module&gt;
    from .lib import xla_client
  File "/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/lib/__init__.py", line 55, in &lt;module&gt;
    from jaxlib import pytree
ImportError: cannot import name 'pytree' from 'jaxlib' (/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jaxlib/__init__.py)
</code></pre></div>
<p dir="auto">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.<br>
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</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="
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 &quot;ilqr_jax_MPC.py&quot;, line 164, in
jac_l, hes_l, jac_l_final, hes_l_final, jac_f = derivative_init()
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/traceback_util.py&quot;, line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/src/api.py&quot;, line 527, in cache_miss
out_flat = xla.xla_call(
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 1937, in bind
return call_bind(self, fun, *args, **params)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 1953, in call_bind
outs = top_trace.process_call(primitive, fun, tracers, params)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 687, in process_call
return primitive.impl(f, *tracers, **params)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py&quot;, line 208, in _xla_call_impl
compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/linear_util.py&quot;, line 295, in memoized_fun
ans = call(fun, *args)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py&quot;, line 257, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
return func(*args, **kwargs)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/dispatch.py&quot;, line 302, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
return func(*args, **kwargs)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py&quot;, line 2188, in trace_to_jaxpr_final2
jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py&quot;, line 2139, in trace_to_subjaxpr_dynamic2
out_tracers = map(trace.full_raise, ans)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/_src/util.py&quot;, line 47, in safe_map
return list(map(f, *args))
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 415, in full_raise
return self.pure(val)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/interpreters/partial_eval.py&quot;, line 1761, in new_const
aval = raise_to_shaped(get_aval(c), weak_type=dtypes.is_weakly_typed(c))
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 1163, in get_aval
return concrete_aval(x)
File &quot;/home/hj/anaconda3/envs/carla/lib/python3.8/site-packages/jax/core.py&quot;, line 1155, in concrete_aval
raise TypeError(f&quot;Value {repr(x)} with type {type(x)} is not a valid JAX &quot;
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Value &lt;CompiledFunction of &lt;function jacfwd..jacfun at 0x7fe3d73ac670&gt;&gt; with type &lt;class 'jaxlib.xla_extension.CompiledFunction'&gt; is not a valid JAX type

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified."><pre class="notranslate"><code class="notranslate">
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 &lt;CompiledFunction of &lt;function jacfwd..jacfun at 0x7fe3d73ac670&gt;&gt; with type &lt;class 'jaxlib.xla_extension.CompiledFunction'&gt; is not a valid JAX type

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
</code></pre></div>
<p dir="auto">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.<br>
Looking forward to your reply.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax0.1.61 jaxlib 0.1.60+CUDA111</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto">py38 ubuntu18.04 CUDA11.1 cuDNN8.0.5 NVIDIA 3060ti</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="python3  tests/lax_numpy_test.py --num_generated_cases=5"><pre lang="bach" class="notranslate"><code class="notranslate">python3  tests/lax_numpy_test.py --num_generated_cases=5
</code></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">[ RUN      ] LaxBackedNumpyTests.testOpIscomplex_bfloat16[3,4]<br>
Fatal Python error: Segmentation fault</p>
<p dir="auto">Current thread 0x00007f19ad5eb740 (most recent call first):<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 537 in backend_compile<br>
File "/data00/home/xxx/jax/jax/_src/profiler.py", line 206 in wrapper<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 583 in compile_or_get_cached<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 614 in from_xla_computation<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 529 in compile<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 170 in _xla_callable_uncached<br>
File "/data00/home/xxx/jax/jax/linear_util.py", line 272 in memoized_fun<br>
File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 143 in _xla_call_impl<br>
File "/data00/home/xxx/jax/jax/core.py", line 614 in process_call<br>
File "/data00/home/xxx/jax/jax/core.py", line 1720 in call_bind<br>
File "/data00/home/xxx/jax/jax/core.py", line 1708 in bind<br>
File "/data00/home/xxx/jax/jax/_src/api.py", line 435 in cache_miss<br>
File "/data00/home/xxx/jax/jax/_src/traceback_util.py", line 162 in reraise_with_filtered_traceback<br>
File "/data00/home/xxx/jax/jax/_src/test_util.py", line 1072 in _CheckAgainstNumpy<br>
File "tests/lax_numpy_test.py", line 589 in testOp<br>
File "/usr/lib/python3.7/contextlib.py", line 74 in inner<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/parameterized.py", line 314 in bound_param_test<br>
File "/usr/lib/python3.7/unittest/case.py", line 615 in run<br>
File "/usr/lib/python3.7/unittest/case.py", line 663 in <strong>call</strong><br>
File "/usr/lib/python3.7/unittest/suite.py", line 122 in run<br>
File "/usr/lib/python3.7/unittest/suite.py", line 84 in <strong>call</strong><br>
File "/usr/lib/python3.7/unittest/suite.py", line 122 in run<br>
File "/usr/lib/python3.7/unittest/suite.py", line 84 in <strong>call</strong><br>
File "/usr/lib/python3.7/unittest/runner.py", line 176 in run<br>
File "/home/yizhengjiao/.local/lib/python3.7/site-packages/absl/testing/_pretty_print_reporter.py", line 86 in run<br>
File "/usr/lib/python3.7/unittest/main.py", line 271 in runTests<br>
File "/usr/lib/python3.7/unittest/main.py", line 101 in <strong>init</strong><br>
File "/home/yizhengjiao/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2537 in _run_and_get_tests_result<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2569 in run_tests<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2164 in main_function<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/app.py", line 258 in _run_main<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/app.py", line 312 in run<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2166 in _run_in_app<br>
File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2049 in main<br>
File "tests/lax_numpy_test.py", line 6317 in <br>
Segmentation fault (core dumped)</p>