megengine.functional.eye

eye(N, M=None, *, dtype='float32', device=None)[source]

Returns a two-dimensional tensor with ones on the diagonal and zeros elsewhere.

Parameters
  • N (int) – number of rows in the output tesnor.

  • M (Optional[int]) – number of columns in the output tesnor. If None, the default number of columns in the output tesnor is equal tos N.

Keyword Arguments
  • dtype (Tensor.dtype, optional) – output tesnor data type. If None, the output tesnor data type must be the default floating-point data type.

  • device (Tensor.device, optional) – device on which to place the created tensor.

See also

If you want to create a diagonal matrix, see diag.

Return type

Tensor

Returns

a tensor where all elements are equal to zero, except for the diagonal, whose values are equal to one.

Examples

>>> F.eye(3)
Tensor([[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]], device=xpux:0)
>>> F.eye(4, 6)
Tensor([[1. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0.]], device=xpux:0)