<p dir="auto"><code class="notranslate">tuple(DeviceArray)</code> evaluates the DeviceArray (at least when using xla) and returns them as np arrays.</p>
<p dir="auto">Example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="ar = jnp.ones(shape=3, dtype=jnp.float32)    # type: jax.interpreters.xla.DeviceArray,
tuple(ar)"><pre class="notranslate"><span class="pl-s1">ar</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>)    <span class="pl-c"># type: jax.interpreters.xla.DeviceArray,</span>
<span class="pl-en">tuple</span>(<span class="pl-s1">ar</span>)</pre></div>
<p dir="auto">produces a tuple of <code class="notranslate">numpy.float32: (1.0, 1.0, 1.0)</code>, while</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="tuple(ar[i] for i in range(ar.shape[0]))"><pre class="notranslate"><span class="pl-en">tuple</span>(<span class="pl-s1">ar</span>[<span class="pl-s1">i</span>] <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">ar</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">0</span>]))</pre></div>
<p dir="auto">produces a tuple of DeviceArrays.</p>
<p dir="auto">It feels like <code class="notranslate">tuple(ar)</code> returning a tuple of DeviceArrays would make sense. Is there some drawback to this?</p>
<p dir="auto">Furthermore, when jitting vs not-jitting the following:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="def tuple_split(ar):
  return tuple(ar)"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">tuple_split</span>(<span class="pl-s1">ar</span>):
  <span class="pl-k">return</span> <span class="pl-en">tuple</span>(<span class="pl-s1">ar</span>)</pre></div>
<p dir="auto">I do get different return types like so:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="print(&quot;orig:&quot;, tuple_split(jnp.ones(shape=3, dtype=jnp.float32)))
&gt; orig: (1.0, 1.0, 1.0)
print(&quot;jit:&quot;, jit_tuple_split(jnp.ones(shape=3, dtype=jnp.float32)))
&gt; jit: (DeviceArray(1., dtype=float32), DeviceArray(1., dtype=float32), DeviceArray(1., dtype=float32))"><pre class="notranslate"><span class="pl-en">print</span>(<span class="pl-s">"orig:"</span>, <span class="pl-en">tuple_split</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>)))
<span class="pl-c1">&gt;</span> <span class="pl-s1">orig</span>: (<span class="pl-c1">1.0</span>, <span class="pl-c1">1.0</span>, <span class="pl-c1">1.0</span>)
<span class="pl-en">print</span>(<span class="pl-s">"jit:"</span>, <span class="pl-en">jit_tuple_split</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>)))
<span class="pl-c1">&gt;</span> <span class="pl-s1">jit</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>), <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>), <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">which I also didn't expect.</p>
<p dir="auto">xla tuple implementation: <a href="https://github.com/google/jax/blob/master/jax/interpreters/xla.py#L1044">https://github.com/google/jax/blob/master/jax/interpreters/xla.py#L1044</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>