2ace85bfe963b8399cc8ec0cdd859a8e0ee31dbd,test/lazy/_lazy_tensor_test_case.py,LazyTensorTestCase,test_diag,#LazyTensorTestCase#,233

Before Change

        evaluated = self.evaluate_lazy_tensor(lazy_tensor)

        res = lazy_tensor.diag()
        actual = evaluated.diag()
        self.assertEqual(res.size(), lazy_tensor.size()[:-1])
        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)

    def test_inv_quad_log_det(self):

After Change

        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)

        res = lazy_tensor.diag()
        actual = torch.stack([flattened_evaluated[i].diag() for i in range(flattened_evaluated.size(0))])
        actual = actual.view(*lazy_tensor.batch_shape, -1)
        self.assertEqual(res.size(), lazy_tensor.size()[:-1])
        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)

Instances