torch.eq的广播机制兼谈快速生成对角掩码

torch.eq用于判断 两个矩阵是不是逐元素相等，或者和第二个值 相等。

x=torch.arange(6)

y=torch.eq(x,3)

print(x,y)

tensor([0, 1, 2, 3, 4, 5])
tensor([False, False, False,  True, False, False])

-----------------------

核心问题：广播机制，如以下代码：

输出  x=[0,1,2,3,4,5]

x.unsqueeze(-1)

[

[0],

[1],

[2],

[3],

[4],

[5]

]

对x进行转置x.T

[[0,1,2,3,4,5]]

print(torch.eq(x,y))

会输出如下矩阵：

tensor([[ True, False, False, False, False, False],
        [False,  True, False, False, False, False],
        [False, False,  True, False, False, False],
        [False, False, False,  True, False, False],
        [False, False, False, False,  True, False],
        [False, False, False, False, False,  True]])

其实就是一个对角线为true的矩阵，怎么实现的？其实很简单。x和y的维度都不相同，进行广播机制，生成两个 6*6矩阵，

X广播后为

[

[[0],[0],[0],[0],[0],[0]],

[[1],[1],[1],[1],[1],[1]]

...

[[5],[5,[5],[5],[5],[5]]

]

y广播后为：

[

[0,1,2,3,4,5],

[0,1,2,3,4,5],

[0,1,2,3,4,5],

[0,1,2,3,4,5],

[0,1,2,3,4,5],

[0,1,2,3,4,5],

]

这样torch.eq(x,y)后只剩 对角线上是true,其它位置是false了。

这样就生成了一个对角线上为true的掩 码矩阵。当然，也可以用eye

print(torch.eye(6))