Table Of Contents
Table Of Contents

mxnet.ndarray.diag

mxnet.ndarray.diag(data=None, k=_Null, axis1=_Null, axis2=_Null, out=None, name=None, **kwargs)

Extracts a diagonal or constructs a diagonal array.

diag’s behavior depends on the input array dimensions:

  • 1-D arrays: constructs a 2-D array with the input as its diagonal, all other elements are zero.

  • N-D arrays: extracts the diagonals of the sub-arrays with axes specified by axis1 and axis2. The output shape would be decided by removing the axes numbered axis1 and axis2 from the input shape and appending to the result a new axis with the size of the diagonals in question.

    For example, when the input shape is (2, 3, 4, 5), axis1 and axis2 are 0 and 2 respectively and k is 0, the resulting shape would be (3, 5, 2).

Examples:

x = [[1, 2, 3],
     [4, 5, 6]]

diag(x) = [1, 5]

diag(x, k=1) = [2, 6]

diag(x, k=-1) = [4]

x = [1, 2, 3]

diag(x) = [[1, 0, 0],
           [0, 2, 0],
           [0, 0, 3]]

diag(x, k=1) = [[0, 1, 0],
                [0, 0, 2],
                [0, 0, 0]]

diag(x, k=-1) = [[0, 0, 0],
                 [1, 0, 0],
                 [0, 2, 0]]

x = [[[1, 2],
      [3, 4]],

     [[5, 6],
      [7, 8]]]

diag(x) = [[1, 7],
           [2, 8]]

diag(x, k=1) = [[3],
                [4]]

diag(x, axis1=-2, axis2=-1) = [[1, 4],
                               [5, 8]]

Defined in src/operator/tensor/diag_op.cc:L87

Parameters:
  • data (NDArray) – Input ndarray
  • k (int, optional, default='0') – Diagonal in question. The default is 0. Use k>0 for diagonals above the main diagonal, and k<0 for diagonals below the main diagonal. If input has shape (S0 S1) k must be between -S0 and S1
  • axis1 (int, optional, default='0') – The first axis of the sub-arrays of interest. Ignored when the input is a 1-D array.
  • axis2 (int, optional, default='1') – The second axis of the sub-arrays of interest. Ignored when the input is a 1-D array.
  • out (NDArray, optional) – The output NDArray to hold the result.
Returns:

out – The output of this function.

Return type:

NDArray or list of NDArrays