megengine.functional.eye

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

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

参数
  • shape – 期望的输出张量的形状。

  • dtype – 数据类型。默认:None

  • device – 矩阵的计算节点。默认:None

返回

单位矩阵。

例如:

import numpy as np
import megengine.functional as F

out = F.eye(4, 6, dtype=np.float32)
print(out.numpy())

输出:

[[1. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0.]]