Merge pull request jax-ml#25191 from houeland:patch-1

PiperOrigin-RevId: 702007643

Author: Google-ML-Automation

jax/_src/numpy/lax_numpy.py

@@ -9076,7 +9076,7 @@ def matmul(a: ArrayLike, b: ArrayLike, *,
 
   Returns:
     array containing the matrix product of the inputs. Shape is ``a.shape[:-1]``
-    if ``b.ndim == 1``, otherwise the shape is ``(..., M)``, where leading
+    if ``b.ndim == 1``, otherwise the shape is ``(..., K, M)``, where leading
     dimensions of ``a`` and ``b`` are broadcast together.
 
   See Also: