Issue 1:
Please:

Check for duplicate issues.

I am running JAX on a Fedora 35 system, with CUDA 11.6, CuDNN 8.2, Driver version 510.60.02

[I installed CuDNN based on the RHEL8 instructions here, since Fedora 35 doesn't seem to officially get the builds for it]

I installed JAX by following the instructions present in the README.md file. On doing that, and simply running the import command, I get the following error

import jax.numpy as jnp

AttributeError Traceback (most recent call last)
Input In [8], in <cell line: 4>()
2 import nibabel as nib
3 import numpy as np
----> 4 import jax.numpy as jnp
6 class NIIHandler():
8 def __init__(self, TRAIN_DATASET_PATH):

File ~/.local/lib/python3.10/site-packages/jax/__init__.py:58, in <module>
38 from jax._src.config import (
39 config as config,
40 enable_checks as enable_checks,
(...)
55 transfer_guard_device_to_host as transfer_guard_device_to_host,
56 )
57 from .core import eval_context as ensure_compile_time_eval
---> 58 from jax._src.api import (
59 ad, # TODO(phawkins): update users to avoid this.
60 block_until_ready,
61 checkpoint as checkpoint,
62 checkpoint_policies as checkpoint_policies,
63 closure_convert as closure_convert,
64 curry, # TODO(phawkins): update users to avoid this.
65 custom_gradient as custom_gradient,
66 custom_jvp as custom_jvp,
67 custom_vjp as custom_vjp,
68 default_backend as default_backend,
69 device_count as device_count,
70 device_get as device_get,
71 device_put as device_put,
72 device_put_sharded as device_put_sharded,
73 device_put_replicated as device_put_replicated,
74 devices as devices,
75 disable_jit as disable_jit,
76 eval_shape as eval_shape,
77 flatten_fun_nokwargs, # TODO(phawkins): update users to avoid this.
78 float0 as float0,
79 grad as grad,
80 hessian as hessian,
81 host_count as host_count,
82 host_id as host_id,
83 host_ids as host_ids,
84 jacobian as jacobian,
85 jacfwd as jacfwd,
86 jacrev as jacrev,
87 jit as jit,
88 jvp as jvp,
89 local_device_count as local_device_count,
90 local_devices as local_devices,
91 linearize as linearize,
92 linear_transpose as linear_transpose,
93 make_jaxpr as make_jaxpr,
94 mask as mask,
95 named_call as named_call,
96 pmap as pmap,
97 process_count as process_count,
98 process_index as process_index,
99 pxla, # TODO(phawkins): update users to avoid this.
100 remat as remat,
101 shapecheck as shapecheck,
102 ShapedArray as ShapedArray,
103 ShapeDtypeStruct as ShapeDtypeStruct,
104 # TODO(phawkins): hide tree* functions from jax, update callers to use
105 # jax.tree_util.
106 treedef_is_leaf,
107 tree_flatten,
108 tree_leaves,
109 tree_map,
110 tree_multimap,
111 tree_structure,
112 tree_transpose,
113 tree_unflatten,
114 value_and_grad as value_and_grad,
115 vjp as vjp,
116 vmap as vmap,
117 xla, # TODO(phawkins): update users to avoid this.
118 xla_computation as xla_computation,
119 )
120 from jax.experimental.maps import soft_pmap as soft_pmap
121 from jax.version import __version__ as __version__

File ~/.local/lib/python3.10/site-packages/jax/_src/api.py:61, in <module>
55 from jax._src import traceback_util
56 from jax._src.api_util import (
57 flatten_fun, apply_flat_fun, flatten_fun_nokwargs, flatten_fun_nokwargs2,
58 argnums_partial, argnums_partial_except, flatten_axes, donation_vector,
59 rebase_donate_argnums, _ensure_index, _ensure_index_tuple,
60 shaped_abstractify, _ensure_str_tuple, argnames_partial_except)
---> 61 from jax._src.lax import lax as lax_internal
62 from jax._src.lib import jax_jit
63 from jax._src.lib import xla_bridge as xb

File ~/.local/lib/python3.10/site-packages/jax/_src/lax/lax.py:1653, in <module>
1651 tan_p = standard_unop(_float | _complex, 'tan')
1652 ad.defjvp2(tan_p, lambda g, ans, x: mul(g, _const(x, 1) + square(ans)))
-> 1653 if jax._src.lib.mlir_api_version >= 11:
1654 mlir.register_lowering(tan_p, partial(_nary_lower_mhlo, chlo.TanOp))
1655 else:

AttributeError: partially initialized module 'jax' has no attribute '_src' (most likely due to a circular import)

I am extremely new to JAX, so please do let me know if there is something else I should be trying instead. Attaching my nvidia-smi and nvcc -- version results below.

Thank you very much!

Issue 2:
Please:

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

https://jax.readthedocs.io/en/latest/jax.random.html#advanced-rng-configuration

In here, the efficiently shardable through jit is only for TPUs. Could you make it for GPU?