Update Returns

jax-ml · Sep 25, 2024 · b9f69f4 · b9f69f4
1 parent 33f77d3
commit b9f69f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -6630,8 +6630,8 @@ def trace(a: ArrayLike, offset: int | ArrayLike = 0, axis1: int = 0, axis2: int
  out: Not used by JAX.
 
  Returns:
- A 0-D array for 2-D input, and in general a N-2 dimensional array for
- N-dimensional input.
+ An array of dimension x.ndim-2 containing the sum of the diagonal elements
+ along axes (axis1, axis2)
 
  See also:
  - :func:`jax.numpy.diag`: Returns the specified diagonal or constructs a diagonal