megengine.functional.eye

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

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

参数

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

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

• device (Optional[CompNode]) – 矩阵的计算节点。默认：None

返回类型

Tensor

返回

单位矩阵。

例如：

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.]]

megengine.functional.linspace megengine.functional.copy