title: IndexError: Array boolean indices must be concrete
body: <p>I'm trying to find the non-zero elements of a list using the LAX-backend implementation of <code>nonzero()</code>.</p>
  <pre><code>from jax import numpy as jnp
  
  Gamma = [[1, 1], [1, 0]]
  print(jnp.nonzero(Gamma[0]))
  </code></pre>
  <p>I'm receiving the error</p>
  <blockquote>
  <p>IndexError: Array boolean indices must be concrete.</p>
  </blockquote>
  <p>However, I can run</p>
  <pre><code>print(np.nonzero(Gamma[0]))
  </code></pre>
  <p>with no issue. A <a href="https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#to-jit-or-not-to-jit" rel="nofollow noreferrer">reference</a> has mentioned that</p>
  <blockquote>
  <p>This [error] is because the function generates an array whose shape is not
  known at compile time: the size of the output depends on the values of
  the input array, and so it is not compatible with JIT.</p>
  </blockquote>
  <p>So how should I fix this issue? Any workarounds?</p>