oneflow.diag

oneflow.diag()

If input is a vector (1-D tensor), then returns a 2-D square tensor with the elements of input as the diagonal. If input is a matrix (2-D tensor), then returns a 1-D tensor with diagonal elements of input.

Parameters

  • input (Tensor) – the input tensor.

  • diagonal (Optional[int], 0) – The diagonal to consider. If diagonal = 0, it is the main diagonal. If diagonal > 0, it is above the main diagonal. If diagonal < 0, it is below the main diagonal. Defaults to 0.

Returns

the output Tensor.

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> arr = np.array(
...     [
...        [1.0, 2.0, 3.0],
...        [4.0, 5.0, 6.0],
...        [7.0, 8.0, 9.0],
...     ]
... )

>>> input = flow.tensor(arr, dtype=flow.float32)
>>> flow.diag(input)
tensor([1., 5., 9.], dtype=oneflow.float32)