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