dq.trace

trace(x: ArrayLike) -> Array

Returns the trace of an array along its last two dimensions.

Parameters

• x (array_like of shape (..., n, n)) –

  Array.

Returns

(array of shape (...)) Trace of x.

Examples

>>> x = jnp.ones((3, 3))
>>> dq.trace(x)
Array(3., dtype=float32)