Пример #1
0
    def test_batch_diag(self):
        lhs = torch.randn(4, 5, 3)
        rhs = torch.randn(4, 3, 5)
        actual = lhs.matmul(rhs)
        actual_diag = torch.cat([
            actual[0].diag().unsqueeze(0),
            actual[1].diag().unsqueeze(0),
            actual[2].diag().unsqueeze(0),
            actual[3].diag().unsqueeze(0),
        ])

        res = MatmulLazyTensor(lhs, rhs)
        self.assertTrue(approx_equal(actual_diag, res.diag()))
Пример #2
0
 def test_diag(self):
     lhs = torch.randn(5, 3)
     rhs = torch.randn(3, 5)
     actual = lhs.matmul(rhs)
     res = MatmulLazyTensor(lhs, rhs)
     self.assertTrue(approx_equal(actual.diag(), res.diag()))