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()))
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()))