megengine.functional.eye¶

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

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

参数

shape – a list, tuple or integer defining the shape of the output tensor.
dtype – the desired data type of the output tensor. Default: float32.
device (Optional[CompNode]) – the desired device of the output tensor. Default: if None, use the default device (see 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