<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>
</ul>
<p dir="auto">I am running <code class="notranslate">JAX</code> on a <code class="notranslate">Fedora 35</code> system, with <code class="notranslate">CUDA 11.6</code>, <code class="notranslate">CuDNN 8.2</code>, <code class="notranslate">Driver version 510.60.02</code></p>
<p dir="auto">[I installed <code class="notranslate">CuDNN</code> based on the <code class="notranslate">RHEL8</code> instructions <a href="https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html" rel="nofollow">here</a>, since <code class="notranslate">Fedora 35</code> doesn't seem to officially get the builds for it]<br>
I installed <code class="notranslate">JAX</code> by following the instructions present in the README.md file. On doing that, and simply running the import command, I get the following error</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="AttributeError                            Traceback (most recent call last)
Input In [8], in &lt;cell line: 4&gt;()
      2 import nibabel as nib
      3 import numpy as np 
----&gt; 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 &lt;module&gt;
     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
---&gt; 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 &lt;module&gt;
     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)
---&gt; 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 &lt;module&gt;
   1651 tan_p = standard_unop(_float | _complex, 'tan')
   1652 ad.defjvp2(tan_p, lambda g, ans, x: mul(g, _const(x, 1) + square(ans)))
-&gt; 1653 if jax._src.lib.mlir_api_version &gt;= 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)"><pre class="notranslate"><span class="pl-v">AttributeError</span>                            <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-v">Input</span> <span class="pl-v">In</span> [<span class="pl-c1">8</span>], <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">cell</span> <span class="pl-s1">line</span>: <span class="pl-c1">4</span><span class="pl-c1">&gt;</span>()
      <span class="pl-c1">2</span> <span class="pl-k">import</span> <span class="pl-s1">nibabel</span> <span class="pl-k">as</span> <span class="pl-s1">nib</span>
      <span class="pl-c1">3</span> <span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span> 
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">4</span> <span class="pl-s1">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
      <span class="pl-c1">6</span> <span class="pl-k">class</span> <span class="pl-v">NIIHandler</span>():
      <span class="pl-c1">8</span>     <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-v">TRAIN_DATASET_PATH</span>):

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

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

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

<span class="pl-v">AttributeError</span>: <span class="pl-s1">partially</span> <span class="pl-s1">initialized</span> <span class="pl-s1">module</span> <span class="pl-s">'jax'</span> <span class="pl-s1">has</span> <span class="pl-s1">no</span> <span class="pl-s1">attribute</span> <span class="pl-s">'_src'</span> (<span class="pl-s1">most</span> <span class="pl-s1">likely</span> <span class="pl-s1">due</span> <span class="pl-s1">to</span> <span class="pl-s1">a</span> <span class="pl-s1">circular</span> <span class="pl-s1">import</span>)</pre></div>
<p dir="auto">I am extremely new to <code class="notranslate">JAX</code>, so please do let me know if there is something else I should be trying instead. Attaching my <code class="notranslate">nvidia-smi</code> and <code class="notranslate">nvcc -- version</code> results below.</p>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/70141886/167195559-3086f698-2dd6-4cd0-8bfc-8d9df4126cb9.png"><img src="https://user-images.githubusercontent.com/70141886/167195559-3086f698-2dd6-4cd0-8bfc-8d9df4126cb9.png" alt="image" style="max-width: 100%;"></a></p>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/70141886/167195599-7fe5c522-fada-47bf-a07a-f7630c76993e.png"><img src="https://user-images.githubusercontent.com/70141886/167195599-7fe5c522-fada-47bf-a07a-f7630c76993e.png" alt="image" style="max-width: 100%;"></a></p>
<p dir="auto">Thank you very much!</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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto"><a href="https://jax.readthedocs.io/en/latest/jax.random.html#advanced-rng-configuration" rel="nofollow">https://jax.readthedocs.io/en/latest/jax.random.html#advanced-rng-configuration</a></p>
<p dir="auto">In here, the efficiently shardable through jit is only for TPUs. Could you make it for GPU?</p>