mlx.core.diagonal#

diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: None | Stream | Device = None) → array#

Return specified diagonals.

If a is 2-D, then a 1-D array containing the diagonal at the given offset is returned.

If a has more than two dimensions, then axis1 and axis2 determine the 2D subarrays from which diagonals are extracted. The new shape is the original shape with axis1 and axis2 removed and a new dimension inserted at the end corresponding to the diagonal.

Parameters:

  • a (array) – Input array

  • offset (int, optional) – Offset of the diagonal from the main diagonal. Can be positive or negative. Default: 0.

  • axis1 (int, optional) – The first axis of the 2-D sub-arrays from which the diagonals should be taken. Default: 0.

  • axis2 (int, optional) – The second axis of the 2-D sub-arrays from which the diagonals should be taken. Default: 1.

Returns:

The diagonals of the array.

Return type:

array