Issue 1:
Description

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.

What jax/jaxlib version are you using?

jax0.1.61 jaxlib 0.1.60+CUDA111

Which accelerator(s) are you using?

GPU

Additional System Info

py38 ubuntu18.04 CUDA11.1 cuDNN8.0.5 NVIDIA 3060ti

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

python3 tests/lax_numpy_test.py --num_generated_cases=5

If applicable, include full error messages/tracebacks.

[ RUN ] LaxBackedNumpyTests.testOpIscomplex_bfloat16[3,4]

Fatal Python error: Segmentation fault

Current thread 0x00007f19ad5eb740 (most recent call first):

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 537 in backend_compile

File "/data00/home/xxx/jax/jax/_src/profiler.py", line 206 in wrapper

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 583 in compile_or_get_cached

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 614 in from_xla_computation

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 529 in compile

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 170 in _xla_callable_uncached

File "/data00/home/xxx/jax/jax/linear_util.py", line 272 in memoized_fun

File "/data00/home/xxx/jax/jax/_src/dispatch.py", line 143 in _xla_call_impl

File "/data00/home/xxx/jax/jax/core.py", line 614 in process_call

File "/data00/home/xxx/jax/jax/core.py", line 1720 in call_bind

File "/data00/home/xxx/jax/jax/core.py", line 1708 in bind

File "/data00/home/xxx/jax/jax/_src/api.py", line 435 in cache_miss

File "/data00/home/xxx/jax/jax/_src/traceback_util.py", line 162 in reraise_with_filtered_traceback

File "/data00/home/xxx/jax/jax/_src/test_util.py", line 1072 in _CheckAgainstNumpy

File "tests/lax_numpy_test.py", line 589 in testOp

File "/usr/lib/python3.7/contextlib.py", line 74 in inner

File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/parameterized.py", line 314 in bound_param_test

File "/usr/lib/python3.7/unittest/case.py", line 615 in run

File "/usr/lib/python3.7/unittest/case.py", line 663 in call

File "/usr/lib/python3.7/unittest/suite.py", line 122 in run

File "/usr/lib/python3.7/unittest/suite.py", line 84 in call

File "/usr/lib/python3.7/unittest/suite.py", line 122 in run

File "/usr/lib/python3.7/unittest/suite.py", line 84 in call

File "/usr/lib/python3.7/unittest/runner.py", line 176 in run

File "/home/yizhengjiao/.local/lib/python3.7/site-packages/absl/testing/_pretty_print_reporter.py", line 86 in run

File "/usr/lib/python3.7/unittest/main.py", line 271 in runTests

File "/usr/lib/python3.7/unittest/main.py", line 101 in init

File "/home/yizhengjiao/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2537 in _run_and_get_tests_result

File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2569 in run_tests

File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2164 in main_function

File "/home/xxx/.local/lib/python3.7/site-packages/absl/app.py", line 258 in _run_main

File "/home/xxx/.local/lib/python3.7/site-packages/absl/app.py", line 312 in run

File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2166 in _run_in_app

File "/home/xxx/.local/lib/python3.7/site-packages/absl/testing/absltest.py", line 2049 in main

File "tests/lax_numpy_test.py", line 6317 in

Segmentation fault (core dumped)