<p dir="auto">Hi, I built jax from source but when importing jax, I get this error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="/usr/local/lib/python3.8/dist-packages/jaxlib/cusolver.py in &lt;module&gt;
     24   from jaxlib import cublas_kernels
     25   for _name, _value in cublas_kernels.registrations().items():
---&gt; 26     xla_client.register_custom_call_target(_name, _value, platform=&quot;CUDA&quot;)
     27 except ImportError:
     28   pass

/usr/local/lib/python3.8/dist-packages/jaxlib/xla_client.py in register_custom_call_target(name, fn, platform)
    432     platform: the target platform.
    433   &quot;&quot;&quot;
--&gt; 434   _xla.register_custom_call_target(name, fn, xla_platform_names[platform])
    435 
    436 

KeyError: 'CUDA'"><pre class="notranslate"><code class="notranslate">/usr/local/lib/python3.8/dist-packages/jaxlib/cusolver.py in &lt;module&gt;
     24   from jaxlib import cublas_kernels
     25   for _name, _value in cublas_kernels.registrations().items():
---&gt; 26     xla_client.register_custom_call_target(_name, _value, platform="CUDA")
     27 except ImportError:
     28   pass

/usr/local/lib/python3.8/dist-packages/jaxlib/xla_client.py in register_custom_call_target(name, fn, platform)
    432     platform: the target platform.
    433   """
--&gt; 434   _xla.register_custom_call_target(name, fn, xla_platform_names[platform])
    435 
    436 

KeyError: 'CUDA'
</code></pre></div>
<p dir="auto">After looking at the source code, it seems that line 26 and 34 in <code class="notranslate">jaxlib/cusolver.py</code> should be:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content=" xla_client.register_custom_call_target(_name, _value, platform=&quot;gpu&quot;)"><pre class="notranslate"> <span class="pl-s1">xla_client</span>.<span class="pl-en">register_custom_call_target</span>(<span class="pl-s1">_name</span>, <span class="pl-s1">_value</span>, <span class="pl-s1">platform</span><span class="pl-c1">=</span><span class="pl-s">"gpu"</span>)</pre></div>
<p dir="auto">to match the xla_platform_names dict in xla_client. When fixed another error is raised:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="  File &quot;/usr/local/lib/python3.8/dist-packages/jaxlib/cuda_prng.py&quot;, line 27, in &lt;module&gt;
    xla_client.register_custom_call_target(_name, _value, platform=&quot;CUDA&quot;)
  File &quot;/usr/local/lib/python3.8/dist-packages/jaxlib/xla_client.py&quot;, line 434, in register_custom_call_target
    _xla.register_custom_call_target(name, fn, xla_platform_names[platform])
KeyError: 'CUDA'"><pre class="notranslate"><code class="notranslate">  File "/usr/local/lib/python3.8/dist-packages/jaxlib/cuda_prng.py", line 27, in &lt;module&gt;
    xla_client.register_custom_call_target(_name, _value, platform="CUDA")
  File "/usr/local/lib/python3.8/dist-packages/jaxlib/xla_client.py", line 434, in register_custom_call_target
    _xla.register_custom_call_target(name, fn, xla_platform_names[platform])
KeyError: 'CUDA'
</code></pre></div>
<p dir="auto">line 434 in <code class="notranslate">jaxlib/xla_client.py</code> should be:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="xla_client.register_custom_call_target(_name, _value, platform=&quot;gpu&quot;)"><pre class="notranslate"><span class="pl-s1">xla_client</span>.<span class="pl-en">register_custom_call_target</span>(<span class="pl-s1">_name</span>, <span class="pl-s1">_value</span>, <span class="pl-s1">platform</span><span class="pl-c1">=</span><span class="pl-s">"gpu"</span>)</pre></div>
<p dir="auto">Let me know if I'm right or if I missed something.</p>
<p dir="auto">Is there any way to create a <code class="notranslate">jax.numpy.array</code> on a specific device, instead of creating it on the default device then moving it via <code class="notranslate">device_put</code>? Also, is there any way to change the default device programmatically? Thanks in advance.</p>