<p dir="auto">I want to sample <code class="notranslate">uint8</code>s uniformly including the maximum value of 255, but this doesn't work:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import random, numpy as jnp
random.randint(random.PRNGKey(0), (10,), 0, 256, jnp.uint8) # [0 0 0 0 0 0 0 0 0 0]"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</span>, <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>), (<span class="pl-c1">10</span>,), <span class="pl-c1">0</span>, <span class="pl-c1">256</span>, <span class="pl-s1">jnp</span>.<span class="pl-s1">uint8</span>) <span class="pl-c"># [0 0 0 0 0 0 0 0 0 0]</span></pre></div>
<p dir="auto">The reason is that maxval is converted into the <code class="notranslate">uint8</code>. This is in contrast to</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import numpy as np
np.random.randint(0, 256, (10,), np.uint8) # [197 133 228 202 216 177 116  78 244 176]"><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">np</span>
<span class="pl-s1">np</span>.<span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-c1">0</span>, <span class="pl-c1">256</span>, (<span class="pl-c1">10</span>,), <span class="pl-s1">np</span>.<span class="pl-s1">uint8</span>) <span class="pl-c"># [197 133 228 202 216 177 116  78 244 176]</span></pre></div>
<p dir="auto">The same happens for all integer datatypes. (Workaround is of course to use a larger type and then cast.)</p>
<p dir="auto">According to randint docstring: <code class="notranslate">Sample uniform random values in [minval, maxval) with given shape/dtype.</code></p>
<p dir="auto">However because of the conversion here </p><div class="Box Box--condensed my-2">
  <div class="Box-header f6">
    <p class="mb-0 text-bold">
      <a href="https://github.com/google/jax/blob/5bbb449ae5849c508194c8eb5b10c101f1fa22ae/jax/_src/random.py#L439">jax/jax/_src/random.py</a>
    </p>
    <p class="mb-0 color-fg-muted">
         Line 439
      in
      <a data-pjax="true" class="commit-tease-sha" href="/google/jax/commit/5bbb449ae5849c508194c8eb5b10c101f1fa22ae">5bbb449</a>
    </p>
  </div>
  <div itemprop="text" class="Box-body p-0 blob-wrapper blob-wrapper-embedded data">
    <table class="highlight tab-size mb-0 js-file-line-container" data-tab-size="8" data-paste-markdown-skip="">

        <tbody><tr class="border-0">
          <td id="L439" class="blob-num border-0 px-3 py-0 color-bg-default" data-line-number="439"></td>
          <td id="LC439" class="blob-code border-0 px-3 py-0 color-bg-default blob-code-inner js-file-line"> <span class="pl-s1">maxval</span> <span class="pl-c1">=</span> <span class="pl-s1">lax</span>.<span class="pl-en">convert_element_type</span>(<span class="pl-s1">maxval</span>, <span class="pl-s1">dtype</span>) </td>
        </tr>
    </tbody></table>
  </div>
</div>
 passing <code class="notranslate">maxval={dtype max value + 1}</code> will generate an all-zero array.<br>
The behavior in numpy is different (and correct):<p></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256, dtype=jnp.uint8)
&gt;&gt; DeviceArray([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=uint8)

np.random.randint(size=(10,), low=0, high=256, dtype=np.uint8)
&gt;&gt; array([ 87, 136, 195, 254,  95,  94, 113,   7,  66,  13], dtype=uint8)"><pre class="notranslate"><code class="notranslate">jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256, dtype=jnp.uint8)
&gt;&gt; DeviceArray([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=uint8)

np.random.randint(size=(10,), low=0, high=256, dtype=np.uint8)
&gt;&gt; array([ 87, 136, 195, 254,  95,  94, 113,   7,  66,  13], dtype=uint8)
</code></pre></div>
<p dir="auto">A workaround is to not pass the dtype to randint (or use a dtype with a higher maximum value), and cast the output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256).astype(jnp.uint8)
&gt;&gt; DeviceArray([178, 233,  99, 238,  39, 149,  46, 139,  68, 142], dtype=uint8)"><pre class="notranslate"><code class="notranslate">jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256).astype(jnp.uint8)
&gt;&gt; DeviceArray([178, 233,  99, 238,  39, 149,  46, 139,  68, 142], dtype=uint8)
</code></pre></div>