megengine.functional.eye

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

返回一个对角线全1而其他元素全0的二维张量。

参数

  • N (int) – 输出张量的行数。

  • M (Optional[int]) – 输出张量的列数。如果为 None, 输出张量列数的默认值为 N.

关键字参数

  • dtype (Tensor.dtype, optional) – 输出张量的数据类型。如果为 None, 则输出张量的默认数据类型必须是浮点数据类型。

  • device (Tensor.device, optional) – 被创建的向量所处的设备

参见

如果你想构造一个对角矩阵，请参考 diag.

返回类型

Tensor

返回

一个张量对角线元素全为1，其他位置元素都为0.

实际案例

>>> F.eye(3)
Tensor([[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]], device=xpux:0)

>>> F.eye(4, 6)
Tensor([[1. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0.]], device=xpux:0)

上一页

megengine.functional.linspace

下一页

megengine.functional.diag