megengine.functional.eye

eye(N, M=None, *, dtype='float32', device=None)[源代码]

返回一个二维张量,其对角线上值均为1,其他位置值为0。

参数
  • N – 一个整数定义行数

  • M – 一个整数定义列数。如果 M 没有被指定,那么列数等于 N 。默认值: None

  • dtype – 输出张量的期望数据类型。默认值:float32

  • device (Optional[CompNode]) – 输出张量的期望设备。默认值:如果为 None,使用全局默认设备(可见 get_default_device)。

返回类型

Tensor

返回

单位矩阵。

实际案例

>>> import numpy as np
>>> out = F.eye(4, 6, dtype=np.float32)
>>> out.numpy()
array([[1., 0., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0., 0.],
       [0., 0., 1., 0., 0., 0.],
       [0., 0., 0., 1., 0., 0.]], dtype=float32)