diagonal

Contents

diagonal#

class brainunit.linalg.diagonal(x, *, offset=0, **kwargs)#

Extract the diagonals of a matrix or stack of matrices.

SaiUnit implementation of numpy.linalg.diagonal().

Unlike saiunit.math.diagonal() (which follows numpy.diagonal() and uses the first two axes by default), this function follows linalg semantics and extracts the diagonal of the last two axes.

Parameters:

  • x (Union[Array, ndarray, number, bool, saiunit.Quantity]) – Input of shape (..., M, N).

  • offset (int) – Offset of the diagonal from the main diagonal (default: 0).

Returns:

out – Diagonals of shape x.shape[:-2] + (K,) with K = min(M, N) for offset=0. Carries the same unit as x.

Return type:

Union[Array, saiunit.Quantity]

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> x = jnp.arange(18.).reshape(2, 3, 3) * u.meter
>>> u.linalg.diagonal(x)
ArrayImpl([[ 0.,  4.,  8.],
           [ 9., 13., 17.]], dtype=float32) * meter
Contents