<h3 dir="auto">Description</h3>
<p dir="auto">Code to reproduce the issue:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.random
import jax.numpy as jnp

def myfun(x, key):
    n = 20
    zeros = jnp.full((n,), 0)
    ones = jnp.full((n,), 1)
    maxval = x[ones] - 1
    return jax.random.randint(key, (n,), zeros, maxval)

key = jax.random.PRNGKey(1)
x = jnp.array([0, 2])
y1 = myfun(x, key)
y2 = jax.jit(myfun)(x, key)
print(y1)
print(y2)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</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">jnp</span>

<span class="pl-k">def</span> <span class="pl-en">myfun</span>(<span class="pl-s1">x</span>, <span class="pl-s1">key</span>):
    <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-c1">20</span>
    <span class="pl-s1">zeros</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">full</span>((<span class="pl-s1">n</span>,), <span class="pl-c1">0</span>)
    <span class="pl-s1">ones</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">full</span>((<span class="pl-s1">n</span>,), <span class="pl-c1">1</span>)
    <span class="pl-s1">maxval</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>[<span class="pl-s1">ones</span>] <span class="pl-c1">-</span> <span class="pl-c1">1</span>
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-s1">key</span>, (<span class="pl-s1">n</span>,), <span class="pl-s1">zeros</span>, <span class="pl-s1">maxval</span>)

<span class="pl-s1">key</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">1</span>)
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">0</span>, <span class="pl-c1">2</span>])
<span class="pl-s1">y1</span> <span class="pl-c1">=</span> <span class="pl-en">myfun</span>(<span class="pl-s1">x</span>, <span class="pl-s1">key</span>)
<span class="pl-s1">y2</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">myfun</span>)(<span class="pl-s1">x</span>, <span class="pl-s1">key</span>)
<span class="pl-en">print</span>(<span class="pl-s1">y1</span>)
<span class="pl-en">print</span>(<span class="pl-s1">y2</span>)</pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 1 1 0 1 1 0 1 1 1 0 0 1 1 0 1 0 1 1 1]"><pre class="notranslate"><code class="notranslate">[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 1 1 0 1 1 0 1 1 1 0 0 1 1 0 1 0 1 1 1]
</code></pre></div>
<p dir="auto">Here <code class="notranslate">maxval</code> is an array full of ones, hence <code class="notranslate">randint</code> should produce zeros only. This is the case with non-jitted version of <code class="notranslate">myfun</code>, while jitted <code class="notranslate">myfun</code> incorrectly produces a mixture of ones and zeros.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax 0.3.16, jaxlib 0.3.15</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto"><em>No response</em></p>
<p dir="auto">I've come across some confusing randint behavior where it seems like maxval is being treated as inclusive rather than exclusive. I've narrowed it down this snippet of code that repros the issue on CPU and GPU:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="@jax.jit
def test():
  batch = jnp.ones([8, 32])

  maxval = batch.shape[1] - jnp.argmax(batch == 1, axis=-1) - 1
  
  seq_crop_first_idx = jax.random.randint(
      jax.random.PRNGKey(0),
      [batch.shape[0]],
      0,
      maxval)
  
  return maxval, seq_crop_first_idx
test()"><pre class="notranslate"><code class="notranslate">@jax.jit
def test():
  batch = jnp.ones([8, 32])

  maxval = batch.shape[1] - jnp.argmax(batch == 1, axis=-1) - 1
  
  seq_crop_first_idx = jax.random.randint(
      jax.random.PRNGKey(0),
      [batch.shape[0]],
      0,
      maxval)
  
  return maxval, seq_crop_first_idx
test()
</code></pre></div>
<p dir="auto">Results in:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(DeviceArray([31, 31, 31, 31, 31, 31, 31, 31], dtype=int32),
 DeviceArray([28,  6, 18,  0, 10,  8, 31,  6], dtype=int32))"><pre class="notranslate"><code class="notranslate">(DeviceArray([31, 31, 31, 31, 31, 31, 31, 31], dtype=int32),
 DeviceArray([28,  6, 18,  0, 10,  8, 31,  6], dtype=int32))
</code></pre></div>
<p dir="auto"><code class="notranslate">maxval</code> is clearly all <code class="notranslate">31</code>, but the second to last output is also <code class="notranslate">31</code>.</p>
<p dir="auto">Here's a colab version: <a href="https://colab.research.google.com/drive/1RJ90xhN-VU7PMwecKhPCCVX1XYhH45FQ?usp=sharing" rel="nofollow">https://colab.research.google.com/drive/1RJ90xhN-VU7PMwecKhPCCVX1XYhH45FQ?usp=sharing</a></p>