Skip to content

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)