<p dir="auto">When we call <code class="notranslate">jax.random.shuffle(key, array, axis=0)</code> on an ndarray of shape (n,m) I'm expecting (based on the docs) that we'll end up with an new array where every <em>row</em> has the same order internally but the order of the rows is shuffled. I'm not super clear that this is guaranteed by the <a href="https://jax.readthedocs.io/en/latest/jax.random.html#jax.random.shuffle" rel="nofollow">docs</a>, which say:</p>
<blockquote>
<p dir="auto">Shuffle the elements of an array uniformly at random along an axis.</p>
</blockquote>
<p dir="auto">But the numpy docs for <a href="https://docs.scipy.org/doc/numpy-1.15.0/reference/generated/numpy.random.shuffle.html" rel="nofollow"><code class="notranslate">numpy.random.shuffle</code></a> do make this super clear, saying that</p>
<blockquote>
<p dir="auto">This function only shuffles the array along the first axis of a multi-dimensional array. The order of sub-arrays is changed but their contents remains the same.</p>
</blockquote>
<p dir="auto">So I'm not super sure the <em>intent</em> is for the Jax behavior to mimic the numpy behavior here, but that is what I was expecting. To demonstrate / reproduce this difference in behavior, see this <a href="https://gist.github.com/zmjjmz/099afcefe9ea3cf3e8015dc8c2e7a340">gist</a></p>
<p dir="auto">Possibly related to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1427923296" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/13038" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/13038/hovercard" href="https://github.com/google/jax/issues/13038">#13038</a>.<br>
Possible duplicate of <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1533775969" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/14022" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/14022/hovercard" href="https://github.com/google/jax/issues/14022">#14022</a>.</p>
<h3 dir="auto">Description</h3>
<p dir="auto">Neither the <code class="notranslate">jax.default_matmul_precision</code> context manager nor the <code class="notranslate">precision=</code> argument to <code class="notranslate">jax.numpy.matmul</code> could change the precision for matrix multiplication on A100, and the result cannot reach a similar precision to that of numpy using float32.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import numpy as np
import jax
import jaxlib
import jax.numpy as jnp
from jax.lax import Precision"><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-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jaxlib</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">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span> <span class="pl-k">import</span> <span class="pl-v">Precision</span></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="A = np.array([[1.1111111, 2.2222222], [3.3333333, 4.4444444]], dtype=np.float32)
Aj = jnp.array(A)"><pre class="notranslate"><span class="pl-v">A</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([[<span class="pl-c1">1.1111111</span>, <span class="pl-c1">2.2222222</span>], [<span class="pl-c1">3.3333333</span>, <span class="pl-c1">4.4444444</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">float32</span>)
<span class="pl-v">Aj</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-v">A</span>)</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# fp32 reference
A @ A.T
&gt;array([[ 6.1728387, 13.580245 ],
&gt;       [13.580245 , 30.864195 ]], dtype=float32)"><pre class="notranslate"><span class="pl-c"># fp32 reference</span>
<span class="pl-v">A</span> @ <span class="pl-v">A</span>.<span class="pl-v">T</span>
<span class="pl-c1">&gt;</span><span class="pl-en">array</span>([[ <span class="pl-c1">6.1728387</span>, <span class="pl-c1">13.580245</span> ],
<span class="pl-c1">&gt;</span>       [<span class="pl-c1">13.580245</span> , <span class="pl-c1">30.864195</span> ]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Aj @ Aj.T
&gt;  Array([[ 6.175251, 13.585552],
&gt;        [13.585552, 30.876255]], dtype=float32)"><pre class="notranslate"><span class="pl-v">Aj</span> @ <span class="pl-v">Aj</span>.<span class="pl-v">T</span>
<span class="pl-c1">&gt;</span>  <span class="pl-v">Array</span>([[ <span class="pl-c1">6.175251</span>, <span class="pl-c1">13.585552</span>],
<span class="pl-c1">&gt;</span>        [<span class="pl-c1">13.585552</span>, <span class="pl-c1">30.876255</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="with jax.default_matmul_precision(&quot;float32&quot;):
    print(Aj @ Aj.T)
&gt;    [[ 6.175251 13.585552]
&gt;     [13.585552 30.876255]]"><pre class="notranslate"><span class="pl-k">with</span> <span class="pl-s1">jax</span>.<span class="pl-en">default_matmul_precision</span>(<span class="pl-s">"float32"</span>):
    <span class="pl-en">print</span>(<span class="pl-v">Aj</span> @ <span class="pl-v">Aj</span>.<span class="pl-v">T</span>)
<span class="pl-c1">&gt;</span>    [[ <span class="pl-c1">6.175251</span> <span class="pl-c1">13.585552</span>]
<span class="pl-c1">&gt;</span>     [<span class="pl-c1">13.585552</span> <span class="pl-c1">30.876255</span>]]</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="jnp.matmul(Aj, Aj.T, precision=Precision.DEFAULT)
&gt;    Array([[ 6.175251, 13.585552],
&gt;           [13.585552, 30.876255]], dtype=float32)"><pre class="notranslate"><span class="pl-s1">jnp</span>.<span class="pl-en">matmul</span>(<span class="pl-v">Aj</span>, <span class="pl-v">Aj</span>.<span class="pl-v">T</span>, <span class="pl-s1">precision</span><span class="pl-c1">=</span><span class="pl-v">Precision</span>.<span class="pl-v">DEFAULT</span>)
<span class="pl-c1">&gt;</span>    <span class="pl-v">Array</span>([[ <span class="pl-c1">6.175251</span>, <span class="pl-c1">13.585552</span>],
<span class="pl-c1">&gt;</span>           [<span class="pl-c1">13.585552</span>, <span class="pl-c1">30.876255</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="jnp.matmul(Aj, Aj.T, precision=Precision.HIGH)
&gt;    Array([[ 6.175251, 13.585552],
&gt;           [13.585552, 30.876255]], dtype=float32)"><pre class="notranslate"><span class="pl-s1">jnp</span>.<span class="pl-en">matmul</span>(<span class="pl-v">Aj</span>, <span class="pl-v">Aj</span>.<span class="pl-v">T</span>, <span class="pl-s1">precision</span><span class="pl-c1">=</span><span class="pl-v">Precision</span>.<span class="pl-v">HIGH</span>)
<span class="pl-c1">&gt;</span>    <span class="pl-v">Array</span>([[ <span class="pl-c1">6.175251</span>, <span class="pl-c1">13.585552</span>],
<span class="pl-c1">&gt;</span>           [<span class="pl-c1">13.585552</span>, <span class="pl-c1">30.876255</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="jnp.matmul(Aj, Aj.T, precision=Precision.HIGHEST)
&gt;    Array([[ 6.175251, 13.585552],
&gt;           [13.585552, 30.876255]], dtype=float32)"><pre class="notranslate"><span class="pl-s1">jnp</span>.<span class="pl-en">matmul</span>(<span class="pl-v">Aj</span>, <span class="pl-v">Aj</span>.<span class="pl-v">T</span>, <span class="pl-s1">precision</span><span class="pl-c1">=</span><span class="pl-v">Precision</span>.<span class="pl-v">HIGHEST</span>)
<span class="pl-c1">&gt;</span>    <span class="pl-v">Array</span>([[ <span class="pl-c1">6.175251</span>, <span class="pl-c1">13.585552</span>],
<span class="pl-c1">&gt;</span>           [<span class="pl-c1">13.585552</span>, <span class="pl-c1">30.876255</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">Development branch</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Python 3.11.1, Ubuntu 20.04</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="    Wed Jan 25 19:41:52 2023       
    +-----------------------------------------------------------------------------+
    | NVIDIA-SMI 470.141.10   Driver Version: 470.141.10   CUDA Version: 11.8     |
    |-------------------------------+----------------------+----------------------+
    | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                               |                      |               MIG M. |
    |===============================+======================+======================|
    |   0  NVIDIA A100-SXM...  On   | 00000000:07:00.0 Off |                    0 |
    | N/A   31C    P0    84W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   1  NVIDIA A100-SXM...  On   | 00000000:0F:00.0 Off |                    0 |
    | N/A   28C    P0    63W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   2  NVIDIA A100-SXM...  On   | 00000000:47:00.0 Off |                    0 |
    | N/A   27C    P0    60W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   3  NVIDIA A100-SXM...  On   | 00000000:4E:00.0 Off |                    0 |
    | N/A   26C    P0    62W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   4  NVIDIA A100-SXM...  On   | 00000000:87:00.0 Off |                    0 |
    | N/A   35C    P0    63W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   5  NVIDIA A100-SXM...  On   | 00000000:90:00.0 Off |                    0 |
    | N/A   32C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   6  NVIDIA A100-SXM...  On   | 00000000:B7:00.0 Off |                    0 |
    | N/A   32C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   7  NVIDIA A100-SXM...  On   | 00000000:BD:00.0 Off |                    0 |
    | N/A   31C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
                                                                                   
    +-----------------------------------------------------------------------------+
    | Processes:                                                                  |
    |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
    |        ID   ID                                                   Usage      |
    |=============================================================================|
    |  No running processes found                                                 |
    +-----------------------------------------------------------------------------+
"><pre class="notranslate"><code class="notranslate">    Wed Jan 25 19:41:52 2023       
    +-----------------------------------------------------------------------------+
    | NVIDIA-SMI 470.141.10   Driver Version: 470.141.10   CUDA Version: 11.8     |
    |-------------------------------+----------------------+----------------------+
    | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                               |                      |               MIG M. |
    |===============================+======================+======================|
    |   0  NVIDIA A100-SXM...  On   | 00000000:07:00.0 Off |                    0 |
    | N/A   31C    P0    84W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   1  NVIDIA A100-SXM...  On   | 00000000:0F:00.0 Off |                    0 |
    | N/A   28C    P0    63W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   2  NVIDIA A100-SXM...  On   | 00000000:47:00.0 Off |                    0 |
    | N/A   27C    P0    60W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   3  NVIDIA A100-SXM...  On   | 00000000:4E:00.0 Off |                    0 |
    | N/A   26C    P0    62W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   4  NVIDIA A100-SXM...  On   | 00000000:87:00.0 Off |                    0 |
    | N/A   35C    P0    63W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   5  NVIDIA A100-SXM...  On   | 00000000:90:00.0 Off |                    0 |
    | N/A   32C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   6  NVIDIA A100-SXM...  On   | 00000000:B7:00.0 Off |                    0 |
    | N/A   32C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
    |   7  NVIDIA A100-SXM...  On   | 00000000:BD:00.0 Off |                    0 |
    | N/A   31C    P0    61W / 400W |      0MiB / 81251MiB |      0%      Default |
    |                               |                      |             Disabled |
    +-------------------------------+----------------------+----------------------+
                                                                                   
    +-----------------------------------------------------------------------------+
    | Processes:                                                                  |
    |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
    |        ID   ID                                                   Usage      |
    |=============================================================================|
    |  No running processes found                                                 |
    +-----------------------------------------------------------------------------+

</code></pre></div>