How to use the gpytorch.test.lazy_tensor_test_case.LazyTensorTestCase function in gpytorch

class TestMatmulLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 1

    def create_lazy_tensor(self):
        lhs = torch.randn(5, 6, requires_grad=True)
        rhs = lhs.clone().detach().transpose(-1, -2)
        covar = MatmulLazyTensor(lhs, rhs)
        return covar

    def evaluate_lazy_tensor(self, lazy_tensor):
        return lazy_tensor.left_lazy_tensor.tensor.matmul(lazy_tensor.right_lazy_tensor.tensor)


class TestMatmulLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 3

    def create_lazy_tensor(self):
        lhs = torch.randn(5, 5, 6, requires_grad=True)
        rhs = lhs.clone().detach().transpose(-1, -2)
        covar = MatmulLazyTensor(lhs, rhs)
        return covar

    def evaluate_lazy_tensor(self, lazy_tensor):
        return lazy_tensor.left_lazy_tensor.tensor.matmul(lazy_tensor.right_lazy_tensor.tensor)


class TestMatmulLazyTensorRectangular(RectangularLazyTensorTestCase, unittest.TestCase):
    def create_lazy_tensor(self):
        lhs = torch.randn(5, 3, requires_grad=True)
        rhs = torch.randn(3, 6, requires_grad=True)

#!/usr/bin/env python3

import unittest

import torch

from gpytorch.lazy import MatmulLazyTensor
from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase, RectangularLazyTensorTestCase


class TestMatmulLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 1

    def create_lazy_tensor(self):
        lhs = torch.randn(5, 6, requires_grad=True)
        rhs = lhs.clone().detach().transpose(-1, -2)
        covar = MatmulLazyTensor(lhs, rhs)
        return covar

    def evaluate_lazy_tensor(self, lazy_tensor):
        return lazy_tensor.left_lazy_tensor.tensor.matmul(lazy_tensor.right_lazy_tensor.tensor)


class TestMatmulLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 3

    def create_lazy_tensor(self):

class TestSumBatchLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 6
    should_test_sample = True

    def create_lazy_tensor(self):
        blocks = torch.randn(12, 4, 4)
        blocks = blocks.transpose(-1, -2).matmul(blocks)
        blocks.requires_grad_(True)
        return SumBatchLazyTensor(NonLazyTensor(blocks))

    def evaluate_lazy_tensor(self, lazy_tensor):
        blocks = lazy_tensor.base_lazy_tensor.tensor
        return blocks.sum(0)


class TestSumBatchLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 6
    should_test_sample = True

    def create_lazy_tensor(self):
        blocks = torch.randn(2, 6, 4, 4)
        blocks = blocks.transpose(-1, -2).matmul(blocks)
        blocks.requires_grad_(True)
        return SumBatchLazyTensor(NonLazyTensor(blocks))

    def evaluate_lazy_tensor(self, lazy_tensor):
        blocks = lazy_tensor.base_lazy_tensor.tensor
        return blocks.view(2, 6, 4, 4).sum(1)


class TestSumBatchLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 6

#!/usr/bin/env python3

import unittest

import torch

from gpytorch.lazy import BlockDiagLazyTensor, NonLazyTensor
from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase


class TestBlockDiagLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 0
    should_test_sample = True

    def create_lazy_tensor(self):
        blocks = torch.randn(8, 4, 4)
        blocks = blocks.matmul(blocks.transpose(-1, -2))
        blocks.add_(torch.eye(4, 4).unsqueeze_(0))
        return BlockDiagLazyTensor(NonLazyTensor(blocks))

    def evaluate_lazy_tensor(self, lazy_tensor):
        blocks = lazy_tensor.base_lazy_tensor.tensor
        actual = torch.zeros(32, 32)
        for i in range(8):
            actual[i * 4 : (i + 1) * 4, i * 4 : (i + 1) * 4] = blocks[i]
        return actual

class TestSumLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 0

    def create_lazy_tensor(self):
        c1 = torch.tensor([[2, 0.5, 0, 0], [5, 1, 2, 0]], dtype=torch.float, requires_grad=True)
        t1 = ToeplitzLazyTensor(c1)
        c2 = torch.tensor([[2, 0.5, 0, 0], [6, 0, 1, -1]], dtype=torch.float, requires_grad=True)
        t2 = ToeplitzLazyTensor(c2)
        return t1 + t2

    def evaluate_lazy_tensor(self, lazy_tensor):
        tensors = [lt.evaluate() for lt in lazy_tensor.lazy_tensors]
        return sum(tensors)


class TestSumLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 0
    # Because these LTs are large, we'll skil the big tests
    skip_slq_tests = True

    def create_lazy_tensor(self):
        c1 = torch.tensor(
            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [5, 1, 2, 0]]], dtype=torch.float, requires_grad=True
        )
        t1 = ToeplitzLazyTensor(c1)
        c2 = torch.tensor(
            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [6, 0, 1, -1]]], dtype=torch.float, requires_grad=True
        )
        t2 = ToeplitzLazyTensor(c2)
        return t1 + t2

    def evaluate_lazy_tensor(self, lazy_tensor):

#!/usr/bin/env python3

import unittest

import torch

from gpytorch.lazy import CatLazyTensor, NonLazyTensor
from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase


class TestCatLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 1

    def create_lazy_tensor(self):
        root = torch.randn(6, 7)
        self.psd_mat = root.matmul(root.t())

        slice1_mat = self.psd_mat[:2, :].requires_grad_()
        slice2_mat = self.psd_mat[2:4, :].requires_grad_()
        slice3_mat = self.psd_mat[4:6, :].requires_grad_()

        slice1 = NonLazyTensor(slice1_mat)
        slice2 = NonLazyTensor(slice2_mat)
        slice3 = NonLazyTensor(slice3_mat)

        return CatLazyTensor(slice1, slice2, slice3, dim=-2)

blocks = torch.randn(8, 4, 4)
        blocks = blocks.matmul(blocks.transpose(-1, -2))
        blocks.add_(torch.eye(4, 4).unsqueeze_(0))
        return BlockInterleavedLazyTensor(NonLazyTensor(blocks))

    def evaluate_lazy_tensor(self, lazy_tensor):
        blocks = lazy_tensor.base_lazy_tensor.tensor
        actual = torch.zeros(32, 32)
        for i in range(8):
            for j in range(4):
                for k in range(4):
                    actual[j * 8 + i, k * 8 + i] = blocks[i, j, k]
        return actual


class TestBlockInterleavedLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 0
    should_test_sample = True

    def create_lazy_tensor(self):
        blocks = torch.randn(2, 6, 4, 4)
        blocks = blocks.matmul(blocks.transpose(-1, -2))
        blocks.add_(torch.eye(4, 4))
        return BlockInterleavedLazyTensor(NonLazyTensor(blocks), block_dim=2)

    def evaluate_lazy_tensor(self, lazy_tensor):
        blocks = lazy_tensor.base_lazy_tensor.tensor
        actual = torch.zeros(2, 24, 24)
        for i in range(2):
            for j in range(6):
                for k in range(4):
                    for l in range(4):

def create_lazy_tensor(self):
        tensor = torch.randn(3, 5, 5)
        tensor = tensor.transpose(-1, -2).matmul(tensor).detach()
        diag = torch.tensor(
            [[1.0, 2.0, 4.0, 2.0, 3.0], [2.0, 1.0, 2.0, 1.0, 4.0], [1.0, 2.0, 2.0, 3.0, 4.0]], requires_grad=True
        )
        return AddedDiagLazyTensor(NonLazyTensor(tensor), DiagLazyTensor(diag))

    def evaluate_lazy_tensor(self, lazy_tensor):
        diag = lazy_tensor._diag_tensor._diag
        tensor = lazy_tensor._lazy_tensor.tensor
        return tensor + torch.diag_embed(diag, dim1=-2, dim2=-1)


class TestAddedDiagLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
    seed = 4
    # Because these LTs are large, we'll skil the big tests
    should_test_sample = False
    skip_slq_tests = True

    def create_lazy_tensor(self):
        tensor = torch.randn(4, 3, 5, 5)
        tensor = tensor.transpose(-1, -2).matmul(tensor).detach()
        diag = (
            torch.tensor(
                [[1.0, 2.0, 4.0, 2.0, 3.0], [2.0, 1.0, 2.0, 1.0, 4.0], [1.0, 2.0, 2.0, 3.0, 4.0]], requires_grad=True
            )
            .repeat(4, 1, 1)
            .detach()
        )
        return AddedDiagLazyTensor(NonLazyTensor(tensor), DiagLazyTensor(diag))

#!/usr/bin/env python3

import unittest

import torch

from gpytorch.lazy import RootLazyTensor
from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase


def make_random_mat(size, rank, batch_shape=torch.Size(())):
    res = torch.randn(*batch_shape, size, rank)
    return res


class TestMulLazyTensor(LazyTensorTestCase, unittest.TestCase):
    seed = 10

    def create_lazy_tensor(self):
        mat1 = make_random_mat(6, 6)
        mat2 = make_random_mat(6, 6)
        res = RootLazyTensor(mat1) * RootLazyTensor(mat2)
        return res.add_diag(torch.tensor(2.0))

    def evaluate_lazy_tensor(self, lazy_tensor):
        diag_tensor = lazy_tensor._diag_tensor.evaluate()
        res = torch.mul(
            lazy_tensor._lazy_tensor.left_lazy_tensor.evaluate(), lazy_tensor._lazy_tensor.right_lazy_tensor.evaluate()
        )
        res = res + diag_tensor
        return res