<p dir="auto">I am seeing this on my local install of jax 0.3.13 with jaxlib 0.3.10</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
r = jnp.ones(3)
r[10]
DeviceArray(1., dtype=float32)"><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>
<span class="pl-s1">r</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>)
<span class="pl-s1">r</span>[<span class="pl-c1">10</span>]
<span class="pl-v">DeviceArray</span>(<span class="pl-c1">1.</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<p dir="auto">Shouldn't <code class="notranslate">r[10]</code> return an out-of-bound <code class="notranslate">IndexError</code> instead?</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>
<h1 dir="auto">GPU/TPU cannot be found after jax process is imported in another open session</h1>
<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.
<ul dir="auto">
<li><a href="https://github.com/google/jax/issues/3004" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/3004/hovercard">This</a> looks different - running 1 piece of code on multiple devices</li>
<li>Concerned <a href="https://github.com/google/jax/issues/5155" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/5155/hovercard">here</a>, unsure if "we haven't published an API for multi-host GPU" is relevant here.</li>
</ul>
</li>
</ul>
<p dir="auto">jax version 0.2.17<br>
jaxlib version 0.1.67</p>
<p dir="auto">Steps to replicate:</p>
<p dir="auto">On a system with multiple devices (in my case, google cloud v2-8 TPU VM).<br>
Open a terminal.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.devices()
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-en">devices</span>()
[<span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,<span class="pl-c1">0</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,<span class="pl-c1">0</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">1</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">2</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,<span class="pl-c1">0</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,<span class="pl-c1">0</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">1</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">4</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,<span class="pl-c1">1</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">5</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>,<span class="pl-c1">1</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">1</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">6</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,<span class="pl-c1">1</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-v">TpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">7</span>, <span class="pl-s1">process_index</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">coords</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,<span class="pl-c1">1</span>,<span class="pl-c1">0</span>), <span class="pl-s1">core_on_chip</span><span class="pl-c1">=</span><span class="pl-c1">1</span>)]</pre></div>
<p dir="auto">With that python session still open, open another terminal/tmux session in my case.<br>
<code class="notranslate">export TF_CPP_MIN_LOG_LEVEL=0</code>:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.devices()
2021-12-03 14:10:23.408083: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:171] XLA service 0x2d12d00 initialized for platform Interpreter (this does not guarantee that XLA will be used). Devices:
2021-12-03 14:10:23.408126: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:179]   StreamExecutor device (0): Interpreter, &lt;undefined&gt;
2021-12-03 14:10:23.419007: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:171] XLA service 0x2cd8ea0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-12-03 14:10:23.419039: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:179]   StreamExecutor device (0): Host, Default Version
2021-12-03 14:10:23.475845: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/tpu_client.cc:222] TPU platform initialization failed: Invalid argument: Invalid fd: -1; Couldn't open device: /dev/accel0 (Operation not permitted); Unable to create Node RegisterInterface for node 0, config: device_path: &quot;/dev/accel0&quot; mode: KERNEL debug_data_directory: &quot;&quot; dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true; could not create driver instance
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-en">devices</span>()
<span class="pl-c1">2021</span><span class="pl-c1">-</span><span class="pl-c1">12</span><span class="pl-c1">-</span><span class="pl-c1">03</span> <span class="pl-c1">14</span>:<span class="pl-c1">10</span>:<span class="pl-c1">23.408083</span>: <span class="pl-v">I</span> <span class="pl-s1">external</span><span class="pl-c1">/</span><span class="pl-s1">org_tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">compiler</span><span class="pl-c1">/</span><span class="pl-s1">xla</span><span class="pl-c1">/</span><span class="pl-s1">service</span><span class="pl-c1">/</span><span class="pl-s1">service</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">171</span>] <span class="pl-v">XLA</span> <span class="pl-s1">service</span> <span class="pl-c1">0x2d12d00</span> <span class="pl-s1">initialized</span> <span class="pl-s1">for</span> <span class="pl-s1">platform</span> <span class="pl-v">Interpreter</span> (<span class="pl-s1">this</span> <span class="pl-s1">does</span> <span class="pl-c1">not</span> <span class="pl-s1">guarantee</span> <span class="pl-s1">that</span> <span class="pl-v">XLA</span> <span class="pl-s1">will</span> <span class="pl-s1">be</span> <span class="pl-s1">used</span>). <span class="pl-v">Devices</span>:
<span class="pl-c1">2021</span><span class="pl-c1">-</span><span class="pl-c1">12</span><span class="pl-c1">-</span><span class="pl-c1">03</span> <span class="pl-c1">14</span>:<span class="pl-c1">10</span>:<span class="pl-c1">23.408126</span>: <span class="pl-v">I</span> <span class="pl-s1">external</span><span class="pl-c1">/</span><span class="pl-s1">org_tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">compiler</span><span class="pl-c1">/</span><span class="pl-s1">xla</span><span class="pl-c1">/</span><span class="pl-s1">service</span><span class="pl-c1">/</span><span class="pl-s1">service</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">179</span>]   <span class="pl-v">StreamExecutor</span> <span class="pl-en">device</span> (<span class="pl-c1">0</span>): <span class="pl-v">Interpreter</span>, <span class="pl-c1">&lt;</span><span class="pl-s1">undefined</span><span class="pl-c1">&gt;</span>
<span class="pl-c1">2021</span><span class="pl-c1">-</span><span class="pl-c1">12</span><span class="pl-c1">-</span><span class="pl-c1">03</span> <span class="pl-c1">14</span>:<span class="pl-c1">10</span>:<span class="pl-c1">23.419007</span>: <span class="pl-v">I</span> <span class="pl-s1">external</span><span class="pl-c1">/</span><span class="pl-s1">org_tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">compiler</span><span class="pl-c1">/</span><span class="pl-s1">xla</span><span class="pl-c1">/</span><span class="pl-s1">service</span><span class="pl-c1">/</span><span class="pl-s1">service</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">171</span>] <span class="pl-v">XLA</span> <span class="pl-s1">service</span> <span class="pl-c1">0x2cd8ea0</span> <span class="pl-s1">initialized</span> <span class="pl-s1">for</span> <span class="pl-s1">platform</span> <span class="pl-v">Host</span> (<span class="pl-s1">this</span> <span class="pl-s1">does</span> <span class="pl-c1">not</span> <span class="pl-s1">guarantee</span> <span class="pl-s1">that</span> <span class="pl-v">XLA</span> <span class="pl-s1">will</span> <span class="pl-s1">be</span> <span class="pl-s1">used</span>). <span class="pl-v">Devices</span>:
<span class="pl-c1">2021</span><span class="pl-c1">-</span><span class="pl-c1">12</span><span class="pl-c1">-</span><span class="pl-c1">03</span> <span class="pl-c1">14</span>:<span class="pl-c1">10</span>:<span class="pl-c1">23.419039</span>: <span class="pl-v">I</span> <span class="pl-s1">external</span><span class="pl-c1">/</span><span class="pl-s1">org_tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">compiler</span><span class="pl-c1">/</span><span class="pl-s1">xla</span><span class="pl-c1">/</span><span class="pl-s1">service</span><span class="pl-c1">/</span><span class="pl-s1">service</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">179</span>]   <span class="pl-v">StreamExecutor</span> <span class="pl-en">device</span> (<span class="pl-c1">0</span>): <span class="pl-v">Host</span>, <span class="pl-v">Default</span> <span class="pl-v">Version</span>
<span class="pl-c1">2021</span><span class="pl-c1">-</span><span class="pl-c1">12</span><span class="pl-c1">-</span><span class="pl-c1">03</span> <span class="pl-c1">14</span>:<span class="pl-c1">10</span>:<span class="pl-c1">23.475845</span>: <span class="pl-v">I</span> <span class="pl-s1">external</span><span class="pl-c1">/</span><span class="pl-s1">org_tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">tensorflow</span><span class="pl-c1">/</span><span class="pl-s1">compiler</span><span class="pl-c1">/</span><span class="pl-s1">xla</span><span class="pl-c1">/</span><span class="pl-s1">pjrt</span><span class="pl-c1">/</span><span class="pl-s1">tpu_client</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">222</span>] <span class="pl-v">TPU</span> <span class="pl-s1">platform</span> <span class="pl-s1">initialization</span> <span class="pl-s1">failed</span>: <span class="pl-v">Invalid</span> <span class="pl-s1">argument</span>: <span class="pl-v">Invalid</span> <span class="pl-s1">fd</span>: <span class="pl-c1">-</span><span class="pl-c1">1</span>; <span class="pl-v">Couldn</span>'<span class="pl-s1">t</span> <span class="pl-s1">open</span> <span class="pl-s1">device</span>: <span class="pl-c1">/</span><span class="pl-s1">dev</span><span class="pl-c1">/</span><span class="pl-en">accel0</span> (<span class="pl-v">Operation</span> <span class="pl-c1">not</span> <span class="pl-s1">permitted</span>); <span class="pl-v">Unable</span> <span class="pl-s1">to</span> <span class="pl-s1">create</span> <span class="pl-v">Node</span> <span class="pl-v">RegisterInterface</span> <span class="pl-s1">for</span> <span class="pl-s1">node</span> <span class="pl-c1">0</span>, <span class="pl-s1">config</span>: <span class="pl-s1">device_path</span>: <span class="pl-s">"/dev/accel0"</span> <span class="pl-s1">mode</span>: <span class="pl-v">KERNEL</span> <span class="pl-s1">debug_data_directory</span>: <span class="pl-s">""</span> <span class="pl-s1">dump_anomalies_only</span>: <span class="pl-s1">true</span> <span class="pl-s1">crash_in_debug_dump</span>: <span class="pl-s1">false</span> <span class="pl-s1">allow_core_dump</span>: <span class="pl-s1">true</span>; <span class="pl-s1">could</span> <span class="pl-c1">not</span> <span class="pl-s1">create</span> <span class="pl-s1">driver</span> <span class="pl-s1">instance</span>
<span class="pl-v">WARNING</span>:<span class="pl-s1">absl</span>:<span class="pl-v">No</span> <span class="pl-v">GPU</span><span class="pl-c1">/</span><span class="pl-v">TPU</span> <span class="pl-s1">found</span>, <span class="pl-s1">falling</span> <span class="pl-s1">back</span> <span class="pl-s1">to</span> <span class="pl-v">CPU</span>. (<span class="pl-v">Set</span> <span class="pl-v">TF_CPP_MIN_LOG_LEVEL</span><span class="pl-c1">=</span><span class="pl-c1">0</span> <span class="pl-c1">and</span> <span class="pl-s1">rerun</span> <span class="pl-s1">for</span> <span class="pl-s1">more</span> <span class="pl-s1">info</span>.)
[<span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)]</pre></div>
<p dir="auto">I expected it to still find all the TPUs, and want to run a script on one TPU at a time.</p>
<p dir="auto">In a more recent version of jax: jax-0.2.25 and jaxlib-0.1.73, the error is more descriptive (which made me realise this is known, possibly not a bug per-se):</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax
WARNING: Logging before InitGoogle() is written to STDERR
I1203 14:15:18.362570  117716 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Not attempting to load libtpu.so in this process.
&gt;&gt;&gt; jax.devices()
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[&lt;jaxlib.xla_extension.Device object at 0x7f042db18570&gt;]"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-v">WARNING</span>: <span class="pl-v">Logging</span> <span class="pl-s1">before</span> <span class="pl-v">InitGoogle</span>() <span class="pl-c1">is</span> <span class="pl-s1">written</span> <span class="pl-s1">to</span> <span class="pl-v">STDERR</span>
<span class="pl-v">I1203</span> <span class="pl-c1">14</span>:<span class="pl-c1">15</span>:<span class="pl-c1">18.362570</span>  <span class="pl-c1">117716</span> <span class="pl-s1">tpu_initializer_helper</span>.<span class="pl-s1">cc</span>:<span class="pl-c1">94</span>] <span class="pl-s1">libtpu</span>.<span class="pl-s1">so</span> <span class="pl-s1">already</span> <span class="pl-c1">in</span> <span class="pl-s1">use</span> <span class="pl-s1">by</span> <span class="pl-s1">another</span> <span class="pl-s1">process</span>. <span class="pl-v">Not</span> <span class="pl-s1">attempting</span> <span class="pl-s1">to</span> <span class="pl-s1">load</span> <span class="pl-s1">libtpu</span>.<span class="pl-s1">so</span> <span class="pl-c1">in</span> <span class="pl-s1">this</span> <span class="pl-s1">process</span>.
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-en">devices</span>()
<span class="pl-v">WARNING</span>:<span class="pl-s1">absl</span>:<span class="pl-v">No</span> <span class="pl-v">GPU</span><span class="pl-c1">/</span><span class="pl-v">TPU</span> <span class="pl-s1">found</span>, <span class="pl-s1">falling</span> <span class="pl-s1">back</span> <span class="pl-s1">to</span> <span class="pl-v">CPU</span>. (<span class="pl-v">Set</span> <span class="pl-v">TF_CPP_MIN_LOG_LEVEL</span><span class="pl-c1">=</span><span class="pl-c1">0</span> <span class="pl-c1">and</span> <span class="pl-s1">rerun</span> <span class="pl-s1">for</span> <span class="pl-s1">more</span> <span class="pl-s1">info</span>.)
[<span class="pl-c1">&lt;</span><span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">Device</span> <span class="pl-s1">object</span> <span class="pl-s1">at</span> <span class="pl-c1">0x7f042db18570</span><span class="pl-c1">&gt;</span>]</pre></div>
<p dir="auto">It seems libtpu.so cannot be used in two separate python processes importing jax, and it would be helpful if this were possible?</p>