oneflow.eye

oneflow.eye(n, m, *, device=None, requires_grad=False, placement=None, sbp) → Tensor

This operator creates a 2-D Tensor with ones on the diagonal and zeros elsewhere.

Parameters

  • n (int) – the number of rows.

  • m (int, optional) – the number of colums with default being n. Defaults to None.

Keyword Arguments

  • device (Union[flow.device, str], optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor.

  • requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

  • placement (oneflow._oneflow_internal.placement, optional) – The placement attribute allows you to specify which physical device the tensor is stored on.

  • sbp (Union[oneflow._oneflow_internal.sbp.sbp, List[oneflow._oneflow_internal.sbp.sbp]], optional) – When creating a global tensor, specify the SBP of the tensor.

Returns

The result tensor with ones on the diagonal and zeros elsewhere.

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow
>>> out = flow.eye(3, 3)
>>> out
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], dtype=oneflow.float32)
>>> out = flow.eye(3, 3, device="cuda")
>>> out
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], device='cuda:0', dtype=oneflow.float32)