<p dir="auto">Hi,</p>
<p dir="auto">When indexing a numpy array with a jax.numpy array, and the jax array has a singleton dimension, the dimension is lost. This behaviour is different if both are numpy arrays, or both are jax arrays.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="array_jnp = jnp.zeros((50, 5))
array_np = np.zeros((50, 5))
idx_jnp = jnp.array([2])
idx_np = np.array([2])
print(array_jnp[idx_jnp])  # prints [[0. 0. 0. 0. 0.]]
print(array_np[idx_jnp])  # prints [0. 0. 0. 0. 0.]    &lt;--------- this is different!
print(array_jnp[idx_np])  # prints [[0. 0. 0. 0. 0.]]
print(array_np[idx_np])  # prints [[0. 0. 0. 0. 0.]]"><pre class="notranslate"><code class="notranslate">array_jnp = jnp.zeros((50, 5))
array_np = np.zeros((50, 5))
idx_jnp = jnp.array([2])
idx_np = np.array([2])
print(array_jnp[idx_jnp])  # prints [[0. 0. 0. 0. 0.]]
print(array_np[idx_jnp])  # prints [0. 0. 0. 0. 0.]    &lt;--------- this is different!
print(array_jnp[idx_np])  # prints [[0. 0. 0. 0. 0.]]
print(array_np[idx_np])  # prints [[0. 0. 0. 0. 0.]]
</code></pre></div>
<p dir="auto">When you try to index a numpy ndarray with a DeviceArray, the numpy array tries to interpret the jax array as a tuple.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import numpy as onp
import jax.numpy as np
x = onp.zeros((5,7))
np_idx = onp.array([1,2,3])
jax_idx = np.array([1,2,3])
x[np_idx]
x[jax_idx] # &lt;- raises IndexError"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">onp</span>
<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-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">5</span>,<span class="pl-c1">7</span>))
<span class="pl-s1">np_idx</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>,<span class="pl-c1">2</span>,<span class="pl-c1">3</span>])
<span class="pl-s1">jax_idx</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>,<span class="pl-c1">2</span>,<span class="pl-c1">3</span>])
<span class="pl-s1">x</span>[<span class="pl-s1">np_idx</span>]
<span class="pl-s1">x</span>[<span class="pl-s1">jax_idx</span>] <span class="pl-c"># &lt;- raises IndexError</span></pre></div>
<p dir="auto">Workaround: put <code class="notranslate">jax_idx</code> in a singleton tuple <code class="notranslate">x[(jax_idx,)]</code></p>
<p dir="auto">This bug resulted in a confusing situation where my function worked when decorated by jax.jit but had a shape mismatch when called on a numpy array.</p>