Issue 1:
When we call jax.random.shuffle(key, array, axis=0) 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 row has the same order internally but the order of the rows is shuffled. I'm not super clear that this is guaranteed by the docs, which say:

Shuffle the elements of an array uniformly at random along an axis.

But the numpy docs for numpy.random.shuffle do make this super clear, saying that

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.

So I'm not super sure the intent 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 gist

Issue 2:
Possibly related to #13038.

Possible duplicate of #14022.

Description

Neither the jax.default_matmul_precision context manager nor the precision= argument to jax.numpy.matmul could change the precision for matrix multiplication on A100, and the result cannot reach a similar precision to that of numpy using float32.

import numpy as np
import jax
import jaxlib
import jax.numpy as jnp
from jax.lax import Precision

A = np.array([[1.1111111, 2.2222222], [3.3333333, 4.4444444]], dtype=np.float32)
Aj = jnp.array(A)

# fp32 reference
A @ A.T
>array([[ 6.1728387, 13.580245 ],
> [13.580245 , 30.864195 ]], dtype=float32)

Aj @ Aj.T
> Array([[ 6.175251, 13.585552],
> [13.585552, 30.876255]], dtype=float32)

with jax.default_matmul_precision("float32"):
print(Aj @ Aj.T)
> [[ 6.175251 13.585552]
> [13.585552 30.876255]]

jnp.matmul(Aj, Aj.T, precision=Precision.DEFAULT)
> Array([[ 6.175251, 13.585552],
> [13.585552, 30.876255]], dtype=float32)

jnp.matmul(Aj, Aj.T, precision=Precision.HIGH)
> Array([[ 6.175251, 13.585552],
> [13.585552, 30.876255]], dtype=float32)

jnp.matmul(Aj, Aj.T, precision=Precision.HIGHEST)
> Array([[ 6.175251, 13.585552],
> [13.585552, 30.876255]], dtype=float32)

What jax/jaxlib version are you using?

Development branch

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.11.1, Ubuntu 20.04

NVIDIA GPU info

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 |
+-----------------------------------------------------------------------------+