fill_diagonal

Contents

fill_diagonal#

class saiunit.math.fill_diagonal(a, val, wrap=False, inplace=False)#

Fill the main diagonal of the given array of any dimensionality.

For an array a with a.ndim >= 2, the diagonal is the list of locations with indices a[i, i, ..., i] all identical.

Parameters:
  • a (Union[saiunit.Quantity, Array, ndarray, number, bool]) – Array in which to fill the diagonal.

  • val (Union[saiunit.Quantity, Array, ndarray, number, bool]) – Value to be written on the diagonal. Its unit must be compatible with that of a.

  • wrap (Optional[bool]) – If True, the diagonal is “wrapped” after a.shape[1] and continues in the first column (for tall matrices). Default is False.

  • inplace (Optional[bool]) – If True, the diagonal is filled in-place. Default is False.

Returns:

out – The input array with the diagonal filled.

Return type:

Union[saiunit.Quantity, Array]

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> u.math.fill_diagonal(jnp.zeros((3, 3)), 5.0)
Array([[5., 0., 0.],
       [0., 5., 0.],
       [0., 0., 5.]], dtype=float32)