Issue 1:
I'm getting a weird error when trying to jit things:

Traceback (most recent call last):
File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 228, in <module>
test_make_batcher()
File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 87, in test_make_batcher
for fn in [make_batcher(5, 2), jit(make_batcher(5, 2))]:
File "/home/ubuntu/dev/research/lottery/cifar10_convnet_run.py", line 79, in make_batcher
splits = list(jnp.arange(1, num_examples // batch_size + 1) * batch_size)
File "/nix/store/ix119bzn54cgvsm98cidfdb4c8w0v1qc-python3.9-jax-0.3.0/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 6747, in deferring_binary_op
return binary_op(self, other)
RuntimeError: UNKNOWN: PTX JIT compiler library not found
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc(56): 'cuLinkCreate(0, nullptr, nullptr, &link_state)'

What does this mean? What .so am I missing? I can confirm that I have libnvidia-ptxjitcompiler.so, libnvidia-ptxjitcompiler.so.1, and libnvidia-ptxjitcompiler.so.495.44 all on my system.

Please:

Check for duplicate issues.

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

from jax import jit
jit(lambda x: 2*x)

If applicable, include full error messages/tracebacks.

https://gist.github.com/samuela/7245561c382842d70a729caecbd11912

Issue 2:
Check for duplicate issues.

I searched is:issue is:open pip install jaxlib metadata

Provide a complete example of how to reproduce the bug:

pip install --upgrade jax==0.2.0 jaxlib==0.1.55+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

If applicable, include full error messages/tracebacks. Returns:

Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html

Requirement already satisfied: jax==0.2.0 in /home/jb/anaconda3/envs/cuda111/lib/python3.8/site-packages (0.2.0)

Collecting jaxlib==0.1.55+cuda110

Using cached https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl (155.2 MB)

WARNING: Discarding https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl (from https://storage.googleapis.com/jax-releases/jax_releases.html). Requested jaxlib==0.1.55+cuda110 from https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55%2Bcuda110-cp38-none-manylinux2010_x86_64.whl has inconsistent version: filename has '0.1.55+cuda110', but metadata has '0.1.55'

ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.55+cuda110 (from versions: 0.1, 0.1.1, 0.1.32, 0.1.33, 0.1.34, 0.1.35, 0.1.36, 0.1.37, 0.1.38, 0.1.39, 0.1.40, 0.1.41, 0.1.42, 0.1.43, 0.1.44, 0.1.45, 0.1.46, 0.1.47, 0.1.48, 0.1.50, 0.1.51, 0.1.52, 0.1.52+cuda100, 0.1.52+cuda101, 0.1.52+cuda102, 0.1.52+cuda110, 0.1.55, 0.1.55+cuda100, 0.1.55+cuda101, 0.1.55+cuda102, 0.1.55+cuda110, 0.1.56, 0.1.56+cuda100, 0.1.56+cuda101, 0.1.56+cuda102, 0.1.56+cuda110, 0.1.56+cuda111, 0.1.57, 0.1.57+cuda101, 0.1.57+cuda102, 0.1.57+cuda110, 0.1.57+cuda111, 0.1.58, 0.1.58+cuda101, 0.1.58+cuda102, 0.1.58+cuda110, 0.1.58+cuda111, 0.1.59, 0.1.59+cuda101, 0.1.59+cuda102, 0.1.59+cuda110, 0.1.59+cuda111, 0.1.60, 0.1.60+cuda101, 0.1.60+cuda102, 0.1.60+cuda110, 0.1.60+cuda111, 0.1.61, 0.1.61+cuda101, 0.1.61+cuda102, 0.1.61+cuda110, 0.1.61+cuda111, 0.1.61+cuda112, 0.1.62, 0.1.62+cuda101, 0.1.62+cuda102, 0.1.62+cuda110, 0.1.62+cuda111, 0.1.62+cuda112, 0.1.63, 0.1.63+cuda101, 0.1.63+cuda102, 0.1.63+cuda110, 0.1.63+cuda111, 0.1.63+cuda112, 0.1.64, 0.1.64+cuda101, 0.1.64+cuda102, 0.1.64+cuda110, 0.1.64+cuda111, 0.1.64+cuda112, 0.1.65, 0.1.65+cuda101, 0.1.65+cuda102, 0.1.65+cuda110, 0.1.65+cuda111, 0.1.65+cuda112, 0.1.66, 0.1.66+cuda101, 0.1.66+cuda102, 0.1.66+cuda110, 0.1.66+cuda111, 0.1.67, 0.1.67+cuda101, 0.1.67+cuda102, 0.1.67+cuda110, 0.1.67+cuda111)

ERROR: No matching distribution found for jaxlib==0.1.55+cuda110

Conda env at time of install:

env.txt

Notable: my cuda toollkit version is 11.1 (to use with torch). I am pinned back on jax 0.1.55 due to DeepMind GLN implementation, if curious.

Edit: Using GTX 1070