a7a12d157766b69cf4b1ddbb5fcdacfe485dc6fa,test/kernels/test_matern_kernel.py,TestMaternKernel,test_ard_separate_batch,#TestMaternKernel#,117

Before Change

        // batch_dims + diag
        res = kernel(a, b, batch_dims=(0, 2)).diag()
        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
        self.assertLess(torch.norm(res - actual), 1e-5)


if __name__ == "__main__":

After Change

            [
                [[[0.0, 0.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]],
                [[[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
                [[[0.0, 0.0], [0.0, 0.0]], [[4.0, 4.0], [0.0, 0.0]]],
            ]
        )

        dist.mul_(math.sqrt(5))
        dist = dist.view(3, 2, 2, 2)
        actual = (dist ** 2 / 3 + dist + 1).mul(torch.exp(-dist))
        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
        self.assertLess(torch.norm(res - actual), 1e-5)

        // batch_dims + diag
        res = kernel(a, b, batch_dims=(0, 2)).diag()
        actual = actual.diagonal(dim1=-2, dim2=-1)
        self.assertLess(torch.norm(res - actual), 1e-5)


if __name__ == "__main__":

Instances