<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>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import os
&gt;&gt;&gt; os.environ[&quot;XLA_FLAGS&quot;] = '--xla_force_host_platform_device_count=8'
&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.devices(&quot;cpu&quot;)
[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]"><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">os</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">os</span>.<span class="pl-s1">environ</span>[<span class="pl-s">"XLA_FLAGS"</span>] <span class="pl-c1">=</span> <span class="pl-s">'--xla_force_host_platform_device_count=8'</span>
<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-s">"cpu"</span>)
[<span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">0</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">1</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">2</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">3</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">4</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">5</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">6</span>), <span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">7</span>)]</pre></div>
<p dir="auto">But the following approaches don't work</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import os
&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.devices(&quot;cpu&quot;)
[CpuDevice(id=0)]
&gt;&gt;&gt; os.environ[&quot;XLA_FLAGS&quot;] = &quot;--xla_force_host_platform_device_count=8&quot;
&gt;&gt;&gt; jax.devices(&quot;cpu&quot;)
[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">os</span>
<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-s">"cpu"</span>)
[<span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)]
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">os</span>.<span class="pl-s1">environ</span>[<span class="pl-s">"XLA_FLAGS"</span>] <span class="pl-c1">=</span> <span class="pl-s">"--xla_force_host_platform_device_count=8"</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-s">"cpu"</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>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import os
&gt;&gt;&gt; import jax
&gt;&gt;&gt; from jax.lib import xla_bridge
&gt;&gt;&gt; jax.devices(&quot;cpu&quot;)
[CpuDevice(id=0)]
&gt;&gt;&gt; os.environ[&quot;XLA_FLAGS&quot;] = &quot;--xla_force_host_platform_device_count=8&quot;
&gt;&gt;&gt; xla_bridge.get_backend.cache_clear()
&gt;&gt;&gt; jax.devices(&quot;cpu&quot;)
[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">os</span>
<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-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lib</span> <span class="pl-k">import</span> <span class="pl-s1">xla_bridge</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-s">"cpu"</span>)
[<span class="pl-v">CpuDevice</span>(<span class="pl-s1">id</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)]
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">os</span>.<span class="pl-s1">environ</span>[<span class="pl-s">"XLA_FLAGS"</span>] <span class="pl-c1">=</span> <span class="pl-s">"--xla_force_host_platform_device_count=8"</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">xla_bridge</span>.<span class="pl-s1">get_backend</span>.<span class="pl-en">cache_clear</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-s">"cpu"</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>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">MyPy says that <code class="notranslate">jax.interpreters.xla.DeviceArray</code> has no attribute "shape", is not indexable, etc.  It seems that it <a href="https://github.com/google/jax/blob/master/jax/interpreters/xla.py#L1053">points</a> to <a href="https://github.com/tensorflow/tensorflow/blob/0306a4d0d69c482eed81b32785d6cba21624c17c/tensorflow/compiler/xla/python/xla_extension/__init__.pyi#L333"><code class="notranslate">DeviceArrayBase</code></a>—not <code class="notranslate">DeviceArray</code>, which has the annotations.</p>
<p dir="auto">I know that type annotations are a work-in-progress. This is just a placeholder issue :)</p>