<p dir="auto">Example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
import numpy as np
a = jnp.arange(10).reshape((2, 5))

type(a)"><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-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">10</span>).<span class="pl-en">reshape</span>((<span class="pl-c1">2</span>, <span class="pl-c1">5</span>))

<span class="pl-en">type</span>(<span class="pl-s1">a</span>)</pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jaxlib.xla_extension.DeviceArray"><pre class="notranslate"><code class="notranslate">jaxlib.xla_extension.DeviceArray
</code></pre></div>
<p dir="auto">But if we do:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="b, c = a"><pre class="notranslate"><code class="notranslate">b, c = a
</code></pre></div>
<p dir="auto">We get</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="type(b), type(c)"><pre class="notranslate"><code class="notranslate">type(b), type(c)
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(numpy.ndarray, numpy.ndarray)"><pre class="notranslate"><code class="notranslate">(numpy.ndarray, numpy.ndarray)
</code></pre></div>
<p dir="auto"><a href="https://colab.research.google.com/gist/romanngg/bc4197d71a404f28224914a2b0dc4481/jnp-np-coversion-when-unpacking.ipynb" rel="nofollow">https://colab.research.google.com/gist/romanngg/bc4197d71a404f28224914a2b0dc4481/jnp-np-coversion-when-unpacking.ipynb</a></p>
<p dir="auto">Looping through a <code class="notranslate">DeviceArray</code> like in a list comprehension, for loop, or enumerate returns <code class="notranslate">numpy.ndarray</code>. Is this expected behaviour?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as np

X = np.array(range(3)).reshape(3,1)
print(type(X))
#&lt;class 'jax.interpreters.xla.DeviceArray'&gt;

x0,x1,x2 = [xi for xi in X]
print(type(x0))
#&lt;type 'numpy.ndarray'&gt;

for xi in X:
  print(type(xi))
#&lt;type 'numpy.ndarray'&gt;
#&lt;type 'numpy.ndarray'&gt;
#&lt;type 'numpy.ndarray'&gt;"><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">np</span>

<span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-en">range</span>(<span class="pl-c1">3</span>)).<span class="pl-en">reshape</span>(<span class="pl-c1">3</span>,<span class="pl-c1">1</span>)
<span class="pl-en">print</span>(<span class="pl-en">type</span>(<span class="pl-v">X</span>))
<span class="pl-c">#&lt;class 'jax.interpreters.xla.DeviceArray'&gt;</span>

<span class="pl-s1">x0</span>,<span class="pl-s1">x1</span>,<span class="pl-s1">x2</span> <span class="pl-c1">=</span> [<span class="pl-s1">xi</span> <span class="pl-k">for</span> <span class="pl-s1">xi</span> <span class="pl-c1">in</span> <span class="pl-v">X</span>]
<span class="pl-en">print</span>(<span class="pl-en">type</span>(<span class="pl-s1">x0</span>))
<span class="pl-c">#&lt;type 'numpy.ndarray'&gt;</span>

<span class="pl-k">for</span> <span class="pl-s1">xi</span> <span class="pl-c1">in</span> <span class="pl-v">X</span>:
  <span class="pl-en">print</span>(<span class="pl-en">type</span>(<span class="pl-s1">xi</span>))
<span class="pl-c">#&lt;type 'numpy.ndarray'&gt;</span>
<span class="pl-c">#&lt;type 'numpy.ndarray'&gt;</span>
<span class="pl-c">#&lt;type 'numpy.ndarray'&gt;</span></pre></div>
<p dir="auto">This is unexpected to me. Why are these not also <code class="notranslate">DeviceArray</code>s?</p>