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
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)