oneflow.diagonal

oneflow.diagonal(input, offset, dim1, dim2)Tensor

Returns a partial view of input with the its diagonal elements with respect to dim1 and dim2 appended as a dimension at the end of the shape.

Parameters
  • input (Tensor) – the input tensor.Must be at least 2-dimensional.

  • offset (Optional[int], 0) – which diagonal to consider. Default: 0 (main diagonal)

  • dim1 (Optional[int], 0) – first dimension with respect to which to take diagonal. Default: 0

  • dim2 (Optional[int], 1) – second dimension with respect to which to take diagonal. Default: 1

Returns

the output Tensor.

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow

>>> input = flow.randn(2,  3,  4)
>>> output = flow.diagonal(input, offset=1, dim1=1, dim2=0)
>>> output.shape
oneflow.Size([4, 1])