jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_simple_sum(self):
        with kernel_arena_scope():
            dtype = torch._C._te.Dtype.Float
            N = 32
            dN = torch._C._te.ExprHandle.int(N)

            A = torch._C._te.Placeholder('A', dtype, [dN])
            B = torch._C._te.Placeholder('B', dtype, [dN])

            def compute(i):
                return A.load([i]) + B.load([i])
            C = torch._C._te.Compute('C', [torch._C._te.DimArg(dN, 'i')], compute)

            loopnest = torch._C._te.LoopNest([C])
            loopnest.prepare_for_codegen()
            stmt = torch._C._te.simplify(loopnest.root_stmt())

            cg = torch._C._te.construct_codegen('ir_eval', stmt, [torch._C._te.BufferArg(x) for x in [A, B, C]])

            tA = torch.rand(N) * 5
            tB = torch.rand(N) * 6
            tC = torch.empty(N)
            cg.call([tA, tB, tC])
            torch.testing.assert_allclose(tA + tB, tC)
//  [python]
def test_simple_sum(self):
        with kernel_arena_scope():
            dtype = torch._C._te.Dtype.Float
            N = 32
            dN = torch._C._te.ExprHandle.int(N)

            A = torch._C._te.Placeholder('A', dtype, [dN])
            B = torch._C._te.Placeholder('B', dtype, [dN])

            def compute(i):
                return A.load([i]) + B.load([i])
            C = torch._C._te.Compute('C', [torch._C._te.DimArg(dN, 'i')], compute)

            loopnest = torch._C._te.LoopNest([C])
            loopnest.prepare_for_codegen()
            stmt = torch._C._te.simplify(loopnest.root_stmt())

            cg = torch._C._te.construct_codegen('ir_eval', stmt, [torch._C._te.BufferArg(x) for x in [A, B, C]])

            tA = torch.rand(N) * 5
            tB = torch.rand(N) * 6
            tC = torch.empty(N)
            cg.call([tA, tB, tC])
            torch.testing.assert_close(tA + tB, tC)
//  [python]
def test_delta(self):
        d = Delta(v=torch.tensor([0.2]))

        def _function(a):
            return 2 * a

        torch.testing.assert_allclose(
            integrate(_function, d, num_samples=10), torch.tensor([0.4])
        )
//  [python]
def test_delta(self):
        d = Delta(v=torch.tensor([0.2]))

        def _function(a):
            return 2 * a

        torch.testing.assert_close(
            integrate(_function, d, num_samples=10), torch.tensor([0.4])
        )
//  [python]
def test_multivariate_normal(self):
        d = MultivariateNormal(torch.tensor([0.2]), scale_tril=1e-6 * torch.eye(1))

        def _function(a):
            return 2 * a

        torch.testing.assert_allclose(
            integrate(_function, d, num_samples=100),
            torch.tensor([0.4]),
            rtol=1e-3,
            atol=1e-3,
        )
//  [python]
def test_multivariate_normal(self):
        d = MultivariateNormal(torch.tensor([0.2]), scale_tril=1e-6 * torch.eye(1))

        def _function(a):
            return 2 * a

        torch.testing.assert_close(
            integrate(_function, d, num_samples=100),
            torch.tensor([0.4]),
            rtol=1e-3,
            atol=1e-3,
        )
//  [python]
def test_scatter_sort(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes = scatter_sort(src, index, descending=True)

    expected = torch.tensor([
        8.3, 2.3, 1.4,
        8.6, 7.6, 6.6, 6.4, 3.7, 1.7,
        9.1, 6.8,
        5.2, 2.0, 1.0, 0.4
    ], device=device)

    torch.testing.assert_allclose(result_values, expected)
    torch.testing.assert_allclose(src[result_indexes], expected)
//  [python]
def test_scatter_sort(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes = scatter_sort(src, index, descending=True)

    expected = torch.tensor([
        8.3, 2.3, 1.4,
        8.6, 7.6, 6.6, 6.4, 3.7, 1.7,
        9.1, 6.8,
        5.2, 2.0, 1.0, 0.4
    ], device=device)

    torch.testing.assert_close(result_values, expected)
    torch.testing.assert_close(src[result_indexes], expected)
//  [python]
def test_correctness(self, omega, batch_size, dim):
        if batch_size:
            x = torch.randn(batch_size, dim)
        else:
            x = torch.randn(dim)

        mm = mellow_max(x, omega)

        torch.testing.assert_allclose(
            mm, 1 / omega * torch.log(torch.mean(torch.exp(omega * x), dim=-1))
        )
        if batch_size:
            torch.testing.assert_allclose(mm.shape, torch.Size([batch_size]))
        else:
            torch.testing.assert_allclose(mm.shape, torch.Size([]))
//  [python]
def test_correctness(self, omega, batch_size, dim):
        if batch_size:
            x = torch.randn(batch_size, dim)
        else:
            x = torch.randn(dim)

        mm = mellow_max(x, omega)

        torch.testing.assert_close(
            mm, 1 / omega * torch.log(torch.mean(torch.exp(omega * x), dim=-1))
        )
        if batch_size:
            torch.testing.assert_close(mm.shape, torch.Size([batch_size]))
        else:
            torch.testing.assert_close(mm.shape, torch.Size([]))
//  [python]
def test_delta(self, batch_size, dim):
        if batch_size:
            mu = torch.randn(batch_size, dim)
            scale = 0 * torch.eye(dim).expand(batch_size, dim, dim)
        else:
            mu = torch.randn(dim)
            scale = 0 * torch.eye(dim)

        d = tensor_to_distribution((mu, scale))
        assert isinstance(d, Delta)
        torch.testing.assert_allclose(d.mean, mu)

        assert d.sample().dtype is torch.get_default_dtype()
        assert d.sample().shape == mu.shape
//  [python]
def test_delta(self, batch_size, dim):
        if batch_size:
            mu = torch.randn(batch_size, dim)
            scale = 0 * torch.eye(dim).expand(batch_size, dim, dim)
        else:
            mu = torch.randn(dim)
            scale = 0 * torch.eye(dim)

        d = tensor_to_distribution((mu, scale))
        assert isinstance(d, Delta)
        torch.testing.assert_close(d.mean, mu)

        assert d.sample().dtype is torch.get_default_dtype()
        assert d.sample().shape == mu.shape
//  [python]
def test_categorical(self, batch_size, dim):
        if batch_size:
            x = torch.randn(batch_size, dim)
        else:
            x = torch.randn(dim)

        d = tensor_to_distribution(x)
        logits = x - x.logsumexp(dim=-1, keepdim=True)
        assert isinstance(d, Categorical)
        torch.testing.assert_allclose(d.logits, logits)
        torch.testing.assert_allclose(d.probs, torch.softmax(x, dim=-1))

        assert isinstance(d.sample(), torch.LongTensor)
        assert d.sample().shape == x.shape[:-1]
//  [python]
def test_categorical(self, batch_size, dim):
        if batch_size:
            x = torch.randn(batch_size, dim)
        else:
            x = torch.randn(dim)

        d = tensor_to_distribution(x)
        logits = x - x.logsumexp(dim=-1, keepdim=True)
        assert isinstance(d, Categorical)
        torch.testing.assert_close(d.logits, logits)
        torch.testing.assert_close(d.probs, torch.softmax(x, dim=-1))

        assert isinstance(d.sample(), torch.LongTensor)
        assert d.sample().shape == x.shape[:-1]
//  [python]
def test_multivariate_normal(self, batch_size, dim):
        if batch_size:
            mu = torch.randn(batch_size, dim)
            scale = 0.1 * torch.eye(dim).expand(batch_size, dim, dim)
        else:
            mu = torch.randn(dim)
            scale = 0.1 * torch.eye(dim)

        d = tensor_to_distribution((mu, scale))
        assert isinstance(d, MultivariateNormal)
        torch.testing.assert_allclose(d.mean, mu)
        torch.testing.assert_allclose(d.covariance_matrix, scale ** 2)

        assert d.sample().dtype is torch.get_default_dtype()
        assert d.sample().shape == mu.shape
//  [python]
def test_multivariate_normal(self, batch_size, dim):
        if batch_size:
            mu = torch.randn(batch_size, dim)
            scale = 0.1 * torch.eye(dim).expand(batch_size, dim, dim)
        else:
            mu = torch.randn(dim)
            scale = 0.1 * torch.eye(dim)

        d = tensor_to_distribution((mu, scale))
        assert isinstance(d, MultivariateNormal)
        torch.testing.assert_close(d.mean, mu)
        torch.testing.assert_close(d.covariance_matrix, scale ** 2)

        assert d.sample().dtype is torch.get_default_dtype()
        assert d.sample().shape == mu.shape
//  [python]
def test_separated_kl(self, dim, batch_size):
        p = self.get_multivariate_normal(dim, batch_size)
        q = self.get_multivariate_normal(dim, batch_size)

        torch.testing.assert_allclose(
            kl_divergence(p, q).sum(), sum(separated_kl(p, q)).sum()
        )

        torch.testing.assert_allclose(
            kl_divergence(q, p).sum(), sum(separated_kl(q, p)).sum()
        )
//  [python]
def test_separated_kl(self, dim, batch_size):
        p = self.get_multivariate_normal(dim, batch_size)
        q = self.get_multivariate_normal(dim, batch_size)

        torch.testing.assert_close(
            kl_divergence(p, q).sum(), sum(separated_kl(p, q)).sum()
        )

        torch.testing.assert_close(
            kl_divergence(q, p).sum(), sum(separated_kl(q, p)).sum()
        )
//  [python]
def test_covariance_matrix(self):
        loc = torch.zeros(self.d)
        exp = (self.W @ torch.diag(self.D) @ self.W.t() + self.s2 * self.Id)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        cov = dist.covariance_matrix
        self.assertEqual(cov.shape, (self.d, self.d))
        tt.assert_allclose(exp, cov)
//  [python]
def test_covariance_matrix(self):
        loc = torch.zeros(self.d)
        exp = (self.W @ torch.diag(self.D) @ self.W.t() + self.s2 * self.Id)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        cov = dist.covariance_matrix
        self.assertEqual(cov.shape, (self.d, self.d))
        torch.testing.assert_close(exp, cov)
//  [python]
def test_precision_matrix(self):
        # tests how accurately the inverse covariance matrix can be computed
        loc = torch.zeros(self.d)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        r = self.W @ torch.diag(self.D) @ self.W.t() + self.s2 * self.Id
        exp = torch.inverse(r)
        tt.assert_allclose(exp, dist.precision_matrix,
                           rtol=1, atol=1 / (math.sqrt(self.d)))
//  [python]
def test_precision_matrix(self):
        # tests how accurately the inverse covariance matrix can be computed
        loc = torch.zeros(self.d)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        r = self.W @ torch.diag(self.D) @ self.W.t() + self.s2 * self.Id
        exp = torch.inverse(r)
        torch.testing.assert_close(exp, dist.precision_matrix,
                           rtol=1, atol=1 / (math.sqrt(self.d)))
//  [python]
def test_log_det(self):
        loc = torch.zeros(self.d)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        cov = dist.covariance_matrix
        res = dist.log_det
        exp = torch.slogdet(cov)[1]
        tt.assert_allclose(res, exp)
//  [python]
def test_log_det(self):
        loc = torch.zeros(self.d)
        dist = MultivariateNormalFactorIdentity(
            loc, self.s2, self.D, self.W)
        cov = dist.covariance_matrix
        res = dist.log_det
        exp = torch.slogdet(cov)[1]
        torch.testing.assert_close(res, exp)
//  [python]
def test_correctness(self, dim, batch_size):
        p = self.get_multivariate_normal(dim, batch_size)
        q = self.get_multivariate_normal(dim, batch_size)
        kl_mean, kl_var = separated_kl(p, q)
        if batch_size is None:
            kl_mean_ = (
                0.5
                * (q.loc - p.loc)
                @ torch.inverse(q.covariance_matrix)
                @ (q.loc - p.loc)
            ).mean()
        else:
            kl_mean_ = (
                0.5
                * (q.loc - p.loc).unsqueeze(1)
                @ torch.inverse(q.covariance_matrix)
                @ (q.loc - p.loc).unsqueeze(-1)
            ).mean()

        n = q.covariance_matrix.shape[-1]
        ratio = torch.inverse(q.covariance_matrix) @ p.covariance_matrix
        kl_var_ = (
            0.5
            * (
                torch.logdet(q.covariance_matrix)
                - torch.logdet(p.covariance_matrix)
                + ratio[..., torch.arange(n), torch.arange(n)].sum(-1)
                - n
            )
        ).mean()

        torch.testing.assert_allclose(kl_mean.mean(), kl_mean_)

        torch.testing.assert_allclose(kl_var.mean(), kl_var_)
//  [python]
def test_correctness(self, dim, batch_size):
        p = self.get_multivariate_normal(dim, batch_size)
        q = self.get_multivariate_normal(dim, batch_size)
        kl_mean, kl_var = separated_kl(p, q)
        if batch_size is None:
            kl_mean_ = (
                0.5
                * (q.loc - p.loc)
                @ torch.inverse(q.covariance_matrix)
                @ (q.loc - p.loc)
            ).mean()
        else:
            kl_mean_ = (
                0.5
                * (q.loc - p.loc).unsqueeze(1)
                @ torch.inverse(q.covariance_matrix)
                @ (q.loc - p.loc).unsqueeze(-1)
            ).mean()

        n = q.covariance_matrix.shape[-1]
        ratio = torch.inverse(q.covariance_matrix) @ p.covariance_matrix
        kl_var_ = (
            0.5
            * (
                torch.logdet(q.covariance_matrix)
                - torch.logdet(p.covariance_matrix)
                + ratio[..., torch.arange(n), torch.arange(n)].sum(-1)
                - n
            )
        ).mean()

        torch.testing.assert_close(kl_mean.mean(), kl_mean_)

        torch.testing.assert_close(kl_var.mean(), kl_var_)
//  [python]
def test_create_index_matrix(self):
        torch.testing.assert_allclose(utils.create_index_matrix(self.nbRow, self.nbCol, self.injunction_table),
                                      self.indexed_matrix)
//  [python]
def test_create_index_matrix(self):
        torch.testing.assert_close(utils.create_index_matrix(self.nbRow, self.nbCol, self.injunction_table),
                                      self.indexed_matrix)
//  [python]
def test_img2mat(self):
        torch.testing.assert_allclose(utils.img2mat(self.image, self.indexed_matrix).type(torch.float),
                                      self.image_matrix)
//  [python]
def test_img2mat(self):
        torch.testing.assert_close(utils.img2mat(self.image, self.indexed_matrix).type(torch.float),
                                      self.image_matrix)
//  [python]
def test_mat2img(self):
        torch.testing.assert_allclose(utils.mat2img(self.image_matrix, self.indexed_matrix),
                                      self.image)
//  [python]
def test_mat2img(self):
        torch.testing.assert_close(utils.mat2img(self.image_matrix, self.indexed_matrix),
                                      self.image)
//  [python]
def test_pool_index_matrix(self):
        torch.testing.assert_allclose(utils.pool_index_matrix(self.indexed_matrix, kernel_type='Hex'),
                                      self.pooled_index_matrix_hex)
        torch.testing.assert_allclose(utils.pool_index_matrix(self.indexed_matrix, kernel_type='Pool'),
                                      self.pooled_index_matrix_pool)
//  [python]
def test_pool_index_matrix(self):
        torch.testing.assert_close(utils.pool_index_matrix(self.indexed_matrix, kernel_type='Hex'),
                                      self.pooled_index_matrix_hex)
        torch.testing.assert_close(utils.pool_index_matrix(self.indexed_matrix, kernel_type='Pool'),
                                      self.pooled_index_matrix_pool)
//  [python]
def test_square_to_hexagonal(self):
        torch.testing.assert_allclose(utils.square_to_hexagonal(self.square_image),
                                      self.square_image_hex)
//  [python]
def test_square_to_hexagonal(self):
        torch.testing.assert_close(utils.square_to_hexagonal(self.square_image),
                                      self.square_image_hex)
//  [python]
def test_build_hexagonal_position(self):
        torch.testing.assert_allclose(utils.build_hexagonal_position(self.square_to_hexagonal_index_matrix),
                                      self.square_to_hexagonal_positions)
//  [python]
def test_build_hexagonal_position(self):
        torch.testing.assert_close(utils.build_hexagonal_position(self.square_to_hexagonal_index_matrix),
                                      self.square_to_hexagonal_positions)
//  [python]
def test_square_to_hexagonal_index_matrix(self):
        torch.testing.assert_allclose(utils.square_to_hexagonal_index_matrix(self.square_image),
                                      self.square_to_hexagonal_index_matrix)
//  [python]
def test_square_to_hexagonal_index_matrix(self):
        torch.testing.assert_close(utils.square_to_hexagonal_index_matrix(self.square_image),
                                      self.square_to_hexagonal_index_matrix)
//  [python]
def test_neighbours_extraction(self):
        torch.testing.assert_allclose(utils.neighbours_extraction(self.indexed_matrix).type(torch.float),
                                      self.neighbours_indices_hex)
        torch.testing.assert_allclose(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Hex', dilation=2).type(torch.float),
            self.neighbours_indices_hex_dilated)
        torch.testing.assert_allclose(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Pool').type(torch.float),
            self.neighbours_indices_pool)
        torch.testing.assert_allclose(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Square', radius=0, stride=2).type(torch.float),
            self.neighbours_indices_one)
        torch.testing.assert_allclose(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Hex', radius=3, retina=True).type(torch.float),
            self.neighbours_indices_retina_hex_3
        )
//  [python]
def test_neighbours_extraction(self):
        torch.testing.assert_close(utils.neighbours_extraction(self.indexed_matrix).type(torch.float),
                                      self.neighbours_indices_hex)
        torch.testing.assert_close(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Hex', dilation=2).type(torch.float),
            self.neighbours_indices_hex_dilated)
        torch.testing.assert_close(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Pool').type(torch.float),
            self.neighbours_indices_pool)
        torch.testing.assert_close(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Square', radius=0, stride=2).type(torch.float),
            self.neighbours_indices_one)
        torch.testing.assert_close(
            utils.neighbours_extraction(self.indexed_matrix, kernel_type='Hex', radius=3, retina=True).type(torch.float),
            self.neighbours_indices_retina_hex_3
        )
//  [python]
def test_build_kernel(self):
        torch.testing.assert_allclose(utils.build_kernel('Hex'),
                                      self.hexagonal_kernel_1)
        torch.testing.assert_allclose(utils.build_kernel('Hex', dilation=2),
                                      self.hexagonal_kernel_1_2)
        torch.testing.assert_allclose(utils.build_kernel('Hex', radius=2),
                                      self.hexagonal_kernel_2)
        torch.testing.assert_allclose(utils.build_kernel('Square'),
                                      self.square_kernel_1)
        torch.testing.assert_allclose(utils.build_kernel('Square', dilation=2),
                                      self.square_kernel_1_2)
        torch.testing.assert_allclose(utils.build_kernel('Square', radius=0),
                                      self.square_kernel_0)
//  [python]
def test_build_kernel(self):
        torch.testing.assert_close(utils.build_kernel('Hex'),
                                      self.hexagonal_kernel_1)
        torch.testing.assert_close(utils.build_kernel('Hex', dilation=2),
                                      self.hexagonal_kernel_1_2)
        torch.testing.assert_close(utils.build_kernel('Hex', radius=2),
                                      self.hexagonal_kernel_2)
        torch.testing.assert_close(utils.build_kernel('Square'),
                                      self.square_kernel_1)
        torch.testing.assert_close(utils.build_kernel('Square', dilation=2),
                                      self.square_kernel_1_2)
        torch.testing.assert_close(utils.build_kernel('Square', radius=0),
                                      self.square_kernel_0)
//  [python]
def test_virtual_sensor_ekf_consistency(generated_data):
    """Check that our Virtual Sensor EKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ekf = torchfilter.filters.VirtualSensorExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ekf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, virtual_sensor_ekf.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance,
        virtual_sensor_ekf.belief_covariance,
    )
//  [python]
def test_virtual_sensor_ekf_consistency(generated_data):
    """Check that our Virtual Sensor EKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ekf = torchfilter.filters.VirtualSensorExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ekf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, virtual_sensor_ekf.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance,
        virtual_sensor_ekf.belief_covariance,
    )
//  [python]
def test_virtual_sensor_eif_consistency(generated_data):
    """Check that our Virtual Sensor EIF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ekf = torchfilter.filters.VirtualSensorExtendedInformationFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ekf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, virtual_sensor_ekf.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance,
        virtual_sensor_ekf.belief_covariance,
    )
//  [python]
def test_virtual_sensor_eif_consistency(generated_data):
    """Check that our Virtual Sensor EIF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ekf = torchfilter.filters.VirtualSensorExtendedInformationFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ekf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, virtual_sensor_ekf.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance,
        virtual_sensor_ekf.belief_covariance,
    )
//  [python]
def test_virtual_sensor_ukf_consistency(generated_data):
    """Check that our Virtual Sensor UKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ukf = torchfilter.filters.VirtualSensorUnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, virtual_sensor_ukf.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance,
        virtual_sensor_ukf.belief_covariance,
        rtol=1e-4,
        atol=5e-4,
    )
//  [python]
def test_virtual_sensor_ukf_consistency(generated_data):
    """Check that our Virtual Sensor UKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_ukf = torchfilter.filters.VirtualSensorUnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        virtual_sensor_model=LinearVirtualSensorModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, virtual_sensor_ukf.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance,
        virtual_sensor_ukf.belief_covariance,
        rtol=1e-4,
        atol=5e-4,
    )
//  [python]
def test_eif_ekf_consistency(generated_data):
    """Check that our EIF and EKF produce consistent results for a linear system. (they
    should be identical)
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    eif = torchfilter.filters.ExtendedInformationFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(eif, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, eif.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance, eif.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_eif_ekf_consistency(generated_data):
    """Check that our EIF and EKF produce consistent results for a linear system. (they
    should be identical)
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    eif = torchfilter.filters.ExtendedInformationFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(eif, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, eif.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance, eif.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_virtual_sensor_srukf_consistency(generated_data):
    """Check that our Virtual Sensor SRUKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_srukf = (
        torchfilter.filters.VirtualSensorSquareRootUnscentedKalmanFilter(
            dynamics_model=LinearDynamicsModel(),
            virtual_sensor_model=LinearVirtualSensorModel(),
        )
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_srukf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, virtual_sensor_srukf.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance,
        virtual_sensor_srukf.belief_covariance,
        rtol=1e-4,
        atol=5e-4,
    )
//  [python]
def test_virtual_sensor_srukf_consistency(generated_data):
    """Check that our Virtual Sensor SRUKF and standard EKF produce consistent results for
    a linear system.
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    virtual_sensor_srukf = (
        torchfilter.filters.VirtualSensorSquareRootUnscentedKalmanFilter(
            dynamics_model=LinearDynamicsModel(),
            virtual_sensor_model=LinearVirtualSensorModel(),
        )
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(virtual_sensor_srukf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, virtual_sensor_srukf.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance,
        virtual_sensor_srukf.belief_covariance,
        rtol=1e-4,
        atol=5e-4,
    )
//  [python]
def test_ukf_ekf_consistency(generated_data):
    """Check that our UKF and EKF produce consistent results for a linear system. (they
    should be identical)
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    ukf = torchfilter.filters.UnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(ekf.belief_mean, ukf.belief_mean)
    torch.testing.assert_allclose(
        ekf.belief_covariance, ukf.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_ukf_ekf_consistency(generated_data):
    """Check that our UKF and EKF produce consistent results for a linear system. (they
    should be identical)
    """
    # Create filters
    ekf = torchfilter.filters.ExtendedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    ukf = torchfilter.filters.UnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(ekf, generated_data)
    _run_filter(ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(ekf.belief_mean, ukf.belief_mean)
    torch.testing.assert_close(
        ekf.belief_covariance, ukf.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_ukf_srukf_consistency(generated_data):
    """Check that our UKF and SRUKF produce consistent results for a linear system.
    (they should be identical)
    """
    # Create filters
    srukf = torchfilter.filters.SquareRootUnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    ukf = torchfilter.filters.UnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(srukf, generated_data)
    _run_filter(ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_allclose(srukf.belief_mean, ukf.belief_mean)
    torch.testing.assert_allclose(
        srukf.belief_covariance, ukf.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_ukf_srukf_consistency(generated_data):
    """Check that our UKF and SRUKF produce consistent results for a linear system.
    (they should be identical)
    """
    # Create filters
    srukf = torchfilter.filters.SquareRootUnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )
    ukf = torchfilter.filters.UnscentedKalmanFilter(
        dynamics_model=LinearDynamicsModel(),
        measurement_model=LinearKalmanFilterMeasurementModel(),
    )

    # Run over data
    _run_filter(srukf, generated_data)
    _run_filter(ukf, generated_data)

    # Check final beliefs
    torch.testing.assert_close(srukf.belief_mean, ukf.belief_mean)
    torch.testing.assert_close(
        srukf.belief_covariance, ukf.belief_covariance, rtol=1e-4, atol=5e-4
    )
//  [python]
def test_unpack_sequences(self):
        X = [torch.Tensor([6, 4, 5]),
             torch.Tensor([1, 4, 5, 7])]
        Y = [torch.Tensor([21, 10, 12, 2, 4, 5]),
             torch.Tensor([1, 4, 11, 13, 14])]
        z, order = pack_sequences(X, Y)
        resX, xlen, resY, ylen = unpack_sequences(z, order)
        tt.assert_allclose(xlen, torch.Tensor([3, 4]).long())
        tt.assert_allclose(ylen, torch.Tensor([6, 5]).long())
        expX = torch.Tensor([[6, 4, 5, 0, 0, 0],
                             [1, 4, 5, 7, 0, 0]])
        expY = torch.Tensor([[21, 10, 12, 2, 4, 5],
                             [1, 4, 11, 13, 14, 0]])
        tt.assert_allclose(expX, resX)
        tt.assert_allclose(expY, resY)
//  [python]
def test_unpack_sequences(self):
        X = [torch.Tensor([6, 4, 5]),
             torch.Tensor([1, 4, 5, 7])]
        Y = [torch.Tensor([21, 10, 12, 2, 4, 5]),
             torch.Tensor([1, 4, 11, 13, 14])]
        z, order = pack_sequences(X, Y)
        resX, xlen, resY, ylen = unpack_sequences(z, order)
        tt.assert_close(xlen, torch.Tensor([3, 4]).long())
        tt.assert_close(ylen, torch.Tensor([6, 5]).long())
        expX = torch.Tensor([[6, 4, 5, 0, 0, 0],
                             [1, 4, 5, 7, 0, 0]])
        expY = torch.Tensor([[21, 10, 12, 2, 4, 5],
                             [1, 4, 11, 13, 14, 0]])
        tt.assert_close(expX, resX)
        tt.assert_close(expY, resY)
//  [python]
def _generate_gradcheck_test(sequencer, fn, poolsize, columnwise, x, xs):
    def _test(self):
        m = ImagePoolingSequencer(
            sequencer=f"{sequencer}-{poolsize}", columnwise=columnwise
        ).to(x.device)
        x.requires_grad_()
        y = m(PaddedTensor(x, xs))
        (dx1,) = torch.autograd.grad(y.data.sum(), (x,))

        for i, (xk, xsk) in enumerate(zip(x, xs.tolist())):
            xk = xk[:, : xsk[0], : xsk[1]].unsqueeze(0).to(x.device)
            xk.requires_grad_()
            yk = fn(
                xk, output_size=(poolsize, xsk[1]) if columnwise else (xsk[0], poolsize)
            )
            (dxk,) = torch.autograd.grad(yk.sum(), (xk,))
            torch.testing.assert_allclose(dxk, (dx1[i : i + 1, :, : xsk[0], : xsk[1]]))

    return _test
//  [python]
def _generate_gradcheck_test(sequencer, fn, poolsize, columnwise, x, xs):
    def _test(self):
        m = ImagePoolingSequencer(
            sequencer=f"{sequencer}-{poolsize}", columnwise=columnwise
        ).to(x.device)
        x.requires_grad_()
        y = m(PaddedTensor(x, xs))
        (dx1,) = torch.autograd.grad(y.data.sum(), (x,))

        for i, (xk, xsk) in enumerate(zip(x, xs.tolist())):
            xk = xk[:, : xsk[0], : xsk[1]].unsqueeze(0).to(x.device)
            xk.requires_grad_()
            yk = fn(
                xk, output_size=(poolsize, xsk[1]) if columnwise else (xsk[0], poolsize)
            )
            (dxk,) = torch.autograd.grad(yk.sum(), (xk,))
            torch.testing.assert_close(dxk, (dx1[i : i + 1, :, : xsk[0], : xsk[1]]))

    return _test
//  [python]
def test_sample_randomness():
    x, x_value, x_probs_and_group = make_discrete_var((3, 2), 17, "x")

    def get_sampler(base_seed: Optional[int] = None):
        return Sampler(RandomSampleSpec((100,), seeder=default_hash_seeder(base_seed)))

    assert get_sampler().sample(x) != get_sampler().sample(x)
    assert torch.abs(get_sampler().sample(x).evaluate() - get_sampler().sample(x).evaluate()).mean() > 0.1

    assert get_sampler(2378).sample(x) == get_sampler(2378).sample(x)
    torch.testing.assert_allclose(get_sampler(2378).sample(x).evaluate(), get_sampler(2378).sample(x).evaluate())

    assert get_sampler().sample(x) != get_sampler().sample(x)
    assert torch.abs(get_sampler().sample(x).evaluate() - get_sampler().sample(x).evaluate()).mean() > 0.1

    def get_manual_sampler(seed: int = 1782, shape: tuple[int, ...] = (100,)):
        torch.manual_seed(seed)
        return Sampler(RandomSampleSpec(shape))

    assert get_manual_sampler().sample(x) == get_manual_sampler().sample(x)
    torch.testing.assert_allclose(get_manual_sampler().sample(x).evaluate(), get_manual_sampler().sample(x).evaluate())
//  [python]
def test_sample_randomness():
    x, x_value, x_probs_and_group = make_discrete_var((3, 2), 17, "x")

    def get_sampler(base_seed: Optional[int] = None):
        return Sampler(RandomSampleSpec((100,), seeder=default_hash_seeder(base_seed)))

    assert get_sampler().sample(x) != get_sampler().sample(x)
    assert torch.abs(get_sampler().sample(x).evaluate() - get_sampler().sample(x).evaluate()).mean() > 0.1

    assert get_sampler(2378).sample(x) == get_sampler(2378).sample(x)
    torch.testing.assert_close(get_sampler(2378).sample(x).evaluate(), get_sampler(2378).sample(x).evaluate())

    assert get_sampler().sample(x) != get_sampler().sample(x)
    assert torch.abs(get_sampler().sample(x).evaluate() - get_sampler().sample(x).evaluate()).mean() > 0.1

    def get_manual_sampler(seed: int = 1782, shape: tuple[int, ...] = (100,)):
        torch.manual_seed(seed)
        return Sampler(RandomSampleSpec(shape))

    assert get_manual_sampler().sample(x) == get_manual_sampler().sample(x)
    torch.testing.assert_close(get_manual_sampler().sample(x).evaluate(), get_manual_sampler().sample(x).evaluate())
//  [python]
def test_multihead_selfattn(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        fake_sample = torch.rand(2, 10, 8)
        fake_sample[1, 3:, :] = 0
        fake_lengths = torch.LongTensor([10, 3])
        padding_mask = lengths_to_padding_mask(fake_lengths)
        fake_sample2 = fake_sample[1:, :3, :]
        padding_mask2 = lengths_to_padding_mask(fake_lengths[1].unsqueeze(0))
        attn = MultiHeadedSelfAttentionModule(8, 4)
        attn.eval()
        attn_out = attn(fake_sample, padding_mask)
        attn_out2 = attn(fake_sample2, padding_mask2)
        torch.testing.assert_allclose(attn_out[1, :3, :], attn_out2[0])
        self.assertTrue(
            torch.all(attn_out[1, 3:, :] == 0.0), f"non-zero entries in {attn_out[1, 3:, :]}")
//  [python]
def test_multihead_selfattn(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        fake_sample = torch.rand(2, 10, 8)
        fake_sample[1, 3:, :] = 0
        fake_lengths = torch.LongTensor([10, 3])
        padding_mask = lengths_to_padding_mask(fake_lengths)
        fake_sample2 = fake_sample[1:, :3, :]
        padding_mask2 = lengths_to_padding_mask(fake_lengths[1].unsqueeze(0))
        attn = MultiHeadedSelfAttentionModule(8, 4)
        attn.eval()
        attn_out = attn(fake_sample, padding_mask)
        attn_out2 = attn(fake_sample2, padding_mask2)
        torch.testing.assert_close(attn_out[1, :3, :], attn_out2[0])
        self.assertTrue(
            torch.all(attn_out[1, 3:, :] == 0.0), f"non-zero entries in {attn_out[1, 3:, :]}")
//  [python]
def test_unresample_nearest_integer_sampling_cpu():

    # Basic Unresample layer
    layer = Unresample('nearest')

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_3x4().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map6(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 5], [9, 11], [3, 10]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    testing.assert_allclose(output, correct_output)
//  [python]
def test_unresample_nearest_integer_sampling_cpu():

    # Basic Unresample layer
    layer = Unresample('nearest')

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_3x4().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map6(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 5], [9, 11], [3, 10]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_unresample_nearest_real_sampling_cpu():

    # Basic Unresample layer
    layer = Unresample('nearest')

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_3x4().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map7(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[1, 6], [10, 0], [0, 11]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    testing.assert_allclose(output, correct_output)
//  [python]
def test_unresample_nearest_real_sampling_cpu():

    # Basic Unresample layer
    layer = Unresample('nearest')

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_3x4().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map7(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[1, 6], [10, 0], [0, 11]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_encoder_batch_unsafe_fails(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        batchnorm_args.input_feat_per_channel = 8
        batchnorm_args.encoder_layers = 3
        batchnorm_args.batch_unsafe_relative_shift = True
        fake_sample = torch.rand(2, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_lengths = torch.LongTensor([27, 13])
        encoder = ConformerEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)
        fake_sample2 = fake_sample[1, :13, :]
        net_out2 = encoder.forward(fake_sample2.unsqueeze(0), fake_lengths[1].unsqueeze(0), return_all_hiddens=True)
        with self.assertRaises(AssertionError) as ae:
            torch.testing.assert_allclose(
                net_out["encoder_out"][0][:4, 1, :],
                net_out2["encoder_out"][0][:, 0, :])
        self.assertTrue("Tensor-likes are not close!" in str(ae.exception))
//  [python]
def test_encoder_batch_unsafe_fails(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        batchnorm_args.input_feat_per_channel = 8
        batchnorm_args.encoder_layers = 3
        batchnorm_args.batch_unsafe_relative_shift = True
        fake_sample = torch.rand(2, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_lengths = torch.LongTensor([27, 13])
        encoder = ConformerEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)
        fake_sample2 = fake_sample[1, :13, :]
        net_out2 = encoder.forward(fake_sample2.unsqueeze(0), fake_lengths[1].unsqueeze(0), return_all_hiddens=True)
        with self.assertRaises(AssertionError) as ae:
            torch.testing.assert_close(
                net_out["encoder_out"][0][:4, 1, :],
                net_out2["encoder_out"][0][:, 0, :])
        self.assertTrue("Tensor-likes are not close!" in str(ae.exception))
//  [python]
def test_augmented_gp(self):
        X_in = torch.from_numpy(self.X_in)
        Y_in = torch.from_numpy(self.Y_in)
        X_out = torch.from_numpy(self.X_out)
        Y_out = torch.from_numpy(self.Y_out)

        augmented_gp = AugmentedLocalGaussianProcess(
            X_in, Y_in, X_out, Y_out, self.gp_model.likelihood, self.kernel
        ).double()
        exact_gp = ExactGaussianProcessModel(X_in, Y_in, self.kernel, self.gp_model.likelihood).double()

        # if augmented_gp.augmented is false, it should behave the same as an exact gp
        output_agp = augmented_gp(X_in)
        output_exact_gp = exact_gp(X_in)
        torch.testing.assert_allclose(output_agp.mean, output_exact_gp.mean)
        torch.testing.assert_allclose(output_agp.covariance_matrix, output_exact_gp.covariance_matrix)

        augmented_gp.eval()
        exact_gp.eval()
        output_agp = augmented_gp(X_out)
        output_exact_gp = exact_gp(X_out)

        torch.testing.assert_allclose(output_agp.mean, output_exact_gp.mean)
        torch.testing.assert_allclose(output_agp.covariance_matrix, output_exact_gp.covariance_matrix)

        # now augmentd_gp is augmented with inducing points, it no longer provides the same output as exact gp
        augmented_gp.set_augment_module(X_inducing=torch.ones([1, self.num_dims]))
        augmented_gp.eval()
        output_agp = augmented_gp(X_out)
        self.assertFalse(torch.equal(output_agp.mean, output_exact_gp.mean))
        self.assertFalse(torch.equal(output_agp.covariance_matrix, output_exact_gp.covariance_matrix))
//  [python]
def test_augmented_gp(self):
        X_in = torch.from_numpy(self.X_in)
        Y_in = torch.from_numpy(self.Y_in)
        X_out = torch.from_numpy(self.X_out)
        Y_out = torch.from_numpy(self.Y_out)

        augmented_gp = AugmentedLocalGaussianProcess(
            X_in, Y_in, X_out, Y_out, self.gp_model.likelihood, self.kernel
        ).double()
        exact_gp = ExactGaussianProcessModel(X_in, Y_in, self.kernel, self.gp_model.likelihood).double()

        # if augmented_gp.augmented is false, it should behave the same as an exact gp
        output_agp = augmented_gp(X_in)
        output_exact_gp = exact_gp(X_in)
        torch.testing.assert_close(output_agp.mean, output_exact_gp.mean)
        torch.testing.assert_close(output_agp.covariance_matrix, output_exact_gp.covariance_matrix)

        augmented_gp.eval()
        exact_gp.eval()
        output_agp = augmented_gp(X_out)
        output_exact_gp = exact_gp(X_out)

        torch.testing.assert_close(output_agp.mean, output_exact_gp.mean)
        torch.testing.assert_close(output_agp.covariance_matrix, output_exact_gp.covariance_matrix)

        # now augmentd_gp is augmented with inducing points, it no longer provides the same output as exact gp
        augmented_gp.set_augment_module(X_inducing=torch.ones([1, self.num_dims]))
        augmented_gp.eval()
        output_agp = augmented_gp(X_out)
        self.assertFalse(torch.equal(output_agp.mean, output_exact_gp.mean))
        self.assertFalse(torch.equal(output_agp.covariance_matrix, output_exact_gp.covariance_matrix))