Jax version:RuntimeError: Internal: Could not find the corresponding function

[songpei1998]

When I ran the command:python -m nerf_sh.eval --chunk 4096 --train_dir $CKPT_ROOT/$SCENE/ --config $CONFIG_FILE --data_dir $DATA_ROOT/$SCENE/:
/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/flax/nn/init.py:35: DeprecationWarning: The flax.nn module is Deprecated, use flax.linen instead. Learn more and find an upgrade guide at https://github.com/google/flax/blob/master/flax/linen/README.md
warnings.warn("The flax.nn module is Deprecated, use flax.linen instead. Learn more and find an upgrade guide at https://github.com/google/flax/blob/master/flax/linen/README.md", DeprecationWarning)
Traceback (most recent call last):
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/runpy.py", line 194, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/runpy.py", line 87, in _run_code
exec(code, run_globals)
File "/home/ubuntu/plenoctree-master/nerf_sh/eval.py", line 137, in
app.run(main)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/home/ubuntu/plenoctree-master/nerf_sh/eval.py", line 46, in main
rng = random.PRNGKey(20200823)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/_src/random.py", line 66, in PRNGKey
k1 = convert(lax.shift_right_logical(seed, lax._const(seed, 32)))
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 383, in shift_right_logical
return shift_right_logical_p.bind(x, y)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 284, in bind
out = top_trace.process_primitive(self, tracers, params)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 622, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/interpreters/xla.py", line 242, in apply_primitive
return compiled_fun(*args)
File "/home/ubuntu/Downloads/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/interpreters/xla.py", line 360, in _execute_compiled_primitive
out_bufs = compiled.execute(input_bufs)
RuntimeError: Internal: Could not find the corresponding function
I also try to upgrade jax version, but the numpy version requirement cannot satisfy the jax and tensorflow.

[songpei1998]

And when I use the jaxlib==0.1.65,it works but with No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)When I use the jaxlib==0.1.65+cuda110,the error is RuntimeError: Internal: Could not find the corresponding function