megengine.functional.eye¶

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

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

参数

shape – 一个表示输出张量的形状的列表，元组或者数字。
dtype – 输出张量的期望数据类型。默认值：float32。
device (Optional[CompNode]) – 输出张量的期望设备。默认值：如果为 None，使用全局默认设备（可见 get_default_device）。

返回类型

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