jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_reim_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_reim_mask = torch.ones((1, fb_dim, 1))
        masked = transforms.apply_real_mask(tf_rep, id_reim_mask, dim=1)
        assert_allclose(masked, tf_rep)
//  [python]
def test_reim_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_reim_mask = torch.ones((1, fb_dim, 1))
        masked = transforms.apply_real_mask(tf_rep, id_reim_mask, dim=1)
        torch.testing.assert_close(masked, tf_rep)
//  [python]
def test_comp_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_complex_mask = torch.cat(
            (torch.ones((1, fb_dim // 2, 1)), torch.zeros((1, fb_dim // 2, 1))), dim=1
        )
        masked = transforms.apply_complex_mask(tf_rep, id_complex_mask, dim=1)
        assert_allclose(masked, tf_rep)
//  [python]
def test_comp_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_complex_mask = torch.cat(
            (torch.ones((1, fb_dim // 2, 1)), torch.zeros((1, fb_dim // 2, 1))), dim=1
        )
        masked = transforms.apply_complex_mask(tf_rep, id_complex_mask, dim=1)
        torch.testing.assert_close(masked, tf_rep)
//  [python]
def test_angle_mag_recompostion(dim):
    """ Test complex --> (mag, angle) --> complex conversions"""
    max_tested_ndim = 4
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    phase = transforms.angle(complex_tensor, dim=dim)
    mag = transforms.mag(complex_tensor, dim=dim)
    tensor_back = transforms.from_magphase(mag, phase, dim=dim)
    assert_allclose(complex_tensor, tensor_back)
//  [python]
def test_angle_mag_recompostion(dim):
    """ Test complex --> (mag, angle) --> complex conversions"""
    max_tested_ndim = 4
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    phase = transforms.angle(complex_tensor, dim=dim)
    mag = transforms.mag(complex_tensor, dim=dim)
    tensor_back = transforms.from_magphase(mag, phase, dim=dim)
    torch.testing.assert_close(complex_tensor, tensor_back)
//  [python]
def test_return_ticket_np_torch(dim):
    """ Test torch --> np --> torch --> np conversion"""
    max_tested_ndim = 4
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    np_array = transforms.to_numpy(complex_tensor, dim=dim)
    tensor_back = transforms.from_numpy(np_array, dim=dim)
    np_back = transforms.to_numpy(tensor_back, dim=dim)
    # Check torch --> np --> torch
    assert_allclose(complex_tensor, tensor_back)
    # Check np --> torch --> np
    np.testing.assert_allclose(np_array, np_back)
//  [python]
def test_return_ticket_np_torch(dim):
    """ Test torch --> np --> torch --> np conversion"""
    max_tested_ndim = 4
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    np_array = transforms.to_numpy(complex_tensor, dim=dim)
    tensor_back = transforms.from_numpy(np_array, dim=dim)
    np_back = transforms.to_numpy(tensor_back, dim=dim)
    # Check torch --> np --> torch
    torch.testing.assert_close(complex_tensor, tensor_back)
    # Check np --> torch --> np
    np.testing.assert_allclose(np_array, np_back)
//  [python]
def test_torchaudio_format(dim, max_tested_ndim):
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    ta_tensor = transforms.to_torchaudio(complex_tensor, dim=dim)
    tensor_back = transforms.from_torchaudio(ta_tensor, dim=dim)
    assert_allclose(complex_tensor, tensor_back)
    assert ta_tensor.shape[-1] == 2
//  [python]
def test_torchaudio_format(dim, max_tested_ndim):
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    ta_tensor = transforms.to_torchaudio(complex_tensor, dim=dim)
    tensor_back = transforms.from_torchaudio(ta_tensor, dim=dim)
    torch.testing.assert_close(complex_tensor, tensor_back)
    assert ta_tensor.shape[-1] == 2
//  [python]
def test_center_freq_correction(kernel_size, stride_factor):
    spec = torch.randn(2, kernel_size + 2, 50)
    stride = None if stride_factor is None else kernel_size // stride_factor
    new_spec = transforms.centerfreq_correction(spec, kernel_size=kernel_size, stride=stride)
    assert spec.shape == new_spec.shape
    assert_allclose(transforms.mag(spec), transforms.mag(new_spec))
//  [python]
def test_center_freq_correction(kernel_size, stride_factor):
    spec = torch.randn(2, kernel_size + 2, 50)
    stride = None if stride_factor is None else kernel_size // stride_factor
    new_spec = transforms.centerfreq_correction(spec, kernel_size=kernel_size, stride=stride)
    assert spec.shape == new_spec.shape
    torch.testing.assert_close(transforms.mag(spec), transforms.mag(new_spec))
//  [python]
def test_superpos_empty(self):
        stack = NoOpStack.empty(1, 3, None)
        policy = torch.ones(1, 3) / 3
        stack.update(policy, NEW_VEC)
        torch.testing.assert_allclose(stack.tapes.tolist(), [[[1/3, 1/3, 0]]])
//  [python]
def test_superpos_empty(self):
        stack = NoOpStack.empty(1, 3, None)
        policy = torch.ones(1, 3) / 3
        stack.update(policy, NEW_VEC)
        torch.testing.assert_close(stack.tapes.tolist(), [[[1/3, 1/3, 0]]])
//  [python]
def test_torch_complex_format(dim, max_tested_ndim):
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    ta_tensor = transforms.to_torch_complex(complex_tensor, dim=dim)
    tensor_back = transforms.from_torch_complex(ta_tensor, dim=dim)
    assert_allclose(complex_tensor, tensor_back)
//  [python]
def test_torch_complex_format(dim, max_tested_ndim):
    # Random tensor shape
    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
    # Make sure complex dimension has even shape
    tensor_shape[dim] = 2 * tensor_shape[dim]
    complex_tensor = torch.randn(tensor_shape)
    ta_tensor = transforms.to_torch_complex(complex_tensor, dim=dim)
    tensor_back = transforms.from_torch_complex(ta_tensor, dim=dim)
    torch.testing.assert_close(complex_tensor, tensor_back)
//  [python]
def test_complex_update(self):
        batch_size = 16
        stack_dim = 10
        policy_vectors = torch.ones(batch_size, 3)
        push_vectors = torch.ones(batch_size, stack_dim)

        stack = NoOpStack.empty(batch_size, stack_dim)
        stack.update(policy_vectors, push_vectors)

        expected = push_vectors.unsqueeze(dim=1)
        torch.testing.assert_allclose(stack.tapes, expected)
//  [python]
def test_complex_update(self):
        batch_size = 16
        stack_dim = 10
        policy_vectors = torch.ones(batch_size, 3)
        push_vectors = torch.ones(batch_size, stack_dim)

        stack = NoOpStack.empty(batch_size, stack_dim)
        stack.update(policy_vectors, push_vectors)

        expected = push_vectors.unsqueeze(dim=1)
        torch.testing.assert_close(stack.tapes, expected)
//  [python]
def _test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self, bucket_view):
        group, group_id, rank = self._init_global_test()
        model = nn.parallel.DistributedDataParallel(ONLY_SBN_NET.npu(rank), device_ids=[rank],
                                                    gradient_as_bucket_view=bucket_view)

        input_var = []
        for i in range(int(WORLD_SIZE)):
            input_var_rank = torch.cat([
                torch.ones(2, 1, 10 ** (i + 1)) * (0.1 ** (i - 1)),
                torch.ones(2, 1, 10 ** (i + 1)) * (0.3 ** (i - 1))
            ], dim=1)
            input_var.append(input_var_rank)

        all_input_var = torch.cat(
            [x.permute(1, 0, 2).contiguous().view(ONLY_SBN_NET.num_features, -1) for x in input_var],
            dim=1
        ).npu(rank)

        for i in range(5):
            y = model(input_var[rank].npu(rank))
            y.mean().backward()

        running_mean, running_var = model.module.running_mean, model.module.running_var
        torch.testing.assert_allclose(running_mean, all_input_var.mean(1))
        torch.testing.assert_allclose(running_var.cpu(), all_input_var.cpu().var(1, unbiased=False))
//  [python]
def _test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self, bucket_view):
        group, group_id, rank = self._init_global_test()
        model = nn.parallel.DistributedDataParallel(ONLY_SBN_NET.npu(rank), device_ids=[rank],
                                                    gradient_as_bucket_view=bucket_view)

        input_var = []
        for i in range(int(WORLD_SIZE)):
            input_var_rank = torch.cat([
                torch.ones(2, 1, 10 ** (i + 1)) * (0.1 ** (i - 1)),
                torch.ones(2, 1, 10 ** (i + 1)) * (0.3 ** (i - 1))
            ], dim=1)
            input_var.append(input_var_rank)

        all_input_var = torch.cat(
            [x.permute(1, 0, 2).contiguous().view(ONLY_SBN_NET.num_features, -1) for x in input_var],
            dim=1
        ).npu(rank)

        for i in range(5):
            y = model(input_var[rank].npu(rank))
            y.mean().backward()

        running_mean, running_var = model.module.running_mean, model.module.running_var
        torch.testing.assert_close(running_mean, all_input_var.mean(1))
        torch.testing.assert_close(running_var.cpu(), all_input_var.cpu().var(1, unbiased=False))
//  [python]
def test_reparametrization(self, out_channels: int, k: int):
        module = MobileOneBlockDown(
            k=k,
            in_channels=3,
            out_channels=out_channels,
        )
        module.eval()

        x = torch.randn(1, 3, 28, 28)

        out1 = module(x)

        rep_module = module.reparametrize()
        rep_module.eval()

        out2 = rep_module(x)

        torch.testing.assert_allclose(out1, out2)
//  [python]
def test_reparametrization(self, out_channels: int, k: int):
        module = MobileOneBlockDown(
            k=k,
            in_channels=3,
            out_channels=out_channels,
        )
        module.eval()

        x = torch.randn(1, 3, 28, 28)

        out1 = module(x)

        rep_module = module.reparametrize()
        rep_module.eval()

        out2 = rep_module(x)

        torch.testing.assert_close(out1, out2)
//  [python]
def test_optimize_for_inference_cpu_torchvision(self):
        models = [
            torchvision.models.resnet18,
            torchvision.models.resnet50,
            torchvision.models.densenet121,
            torchvision.models.shufflenet_v2_x1_0,
            torchvision.models.vgg16,
            torchvision.models.mobilenet_v2,
            torchvision.models.mnasnet1_0,
            torchvision.models.resnext50_32x4d,
        ]
        with torch.no_grad():
            for model_type in models:
                model = model_type()
                C, H, W, = (
                    3,
                    224,
                    224,
                )
                inp = torch.randn(3, C, H, W)
                model(inp)
                model.eval()
                inp = torch.randn(1, C, H, W)
                heuristic = optimization.gen_mkl_autotuner(inp, iters=0, warmup=0)
                optimized_model = optimization.optimize_for_inference(model)

                orig_out = model(inp)
                new_out = optimized_model(inp)
                torch.testing.assert_allclose(orig_out, new_out)
//  [python]
def test_optimize_for_inference_cpu_torchvision(self):
        models = [
            torchvision.models.resnet18,
            torchvision.models.resnet50,
            torchvision.models.densenet121,
            torchvision.models.shufflenet_v2_x1_0,
            torchvision.models.vgg16,
            torchvision.models.mobilenet_v2,
            torchvision.models.mnasnet1_0,
            torchvision.models.resnext50_32x4d,
        ]
        with torch.no_grad():
            for model_type in models:
                model = model_type()
                C, H, W, = (
                    3,
                    224,
                    224,
                )
                inp = torch.randn(3, C, H, W)
                model(inp)
                model.eval()
                inp = torch.randn(1, C, H, W)
                heuristic = optimization.gen_mkl_autotuner(inp, iters=0, warmup=0)
                optimized_model = optimization.optimize_for_inference(model)

                orig_out = model(inp)
                new_out = optimized_model(inp)
                torch.testing.assert_close(orig_out, new_out)
//  [python]
def test_optimize_for_inference_cpu(self):
        import torch.nn as nn

        class Foo(nn.Module):
            def __init__(self):
                super().__init__()
                layers = []
                layers2 = []
                for _ in range(10):
                    layers.append(nn.Conv2d(3, 3, 1))
                    layers.append(nn.BatchNorm2d(3))
                    layers.append(nn.ReLU())

                    layers2.append(nn.Conv2d(3, 3, 1))
                    layers2.append(nn.BatchNorm2d(3))
                    layers2.append(nn.ReLU())
                self.model = nn.Sequential(*layers)
                self.model2 = nn.Sequential(*layers2)

            def forward(self, x):
                return self.model(x) + self.model2(x)

        N, C, H, W, = (
            1,
            3,
            224,
            224,
        )
        inp = torch.randn(N, C, H, W)
        with torch.no_grad():
            model = Foo().eval()
            optimized_model = optimization.optimize_for_inference(model)
            torch.testing.assert_allclose(model(inp), optimized_model(inp))

            optimized_model2 = optimization.optimize_for_inference(
                model, pass_config={"remove_dropout": False}
            )
            torch.testing.assert_allclose(model(inp), optimized_model2(inp))
//  [python]
def test_optimize_for_inference_cpu(self):
        import torch.nn as nn

        class Foo(nn.Module):
            def __init__(self):
                super().__init__()
                layers = []
                layers2 = []
                for _ in range(10):
                    layers.append(nn.Conv2d(3, 3, 1))
                    layers.append(nn.BatchNorm2d(3))
                    layers.append(nn.ReLU())

                    layers2.append(nn.Conv2d(3, 3, 1))
                    layers2.append(nn.BatchNorm2d(3))
                    layers2.append(nn.ReLU())
                self.model = nn.Sequential(*layers)
                self.model2 = nn.Sequential(*layers2)

            def forward(self, x):
                return self.model(x) + self.model2(x)

        N, C, H, W, = (
            1,
            3,
            224,
            224,
        )
        inp = torch.randn(N, C, H, W)
        with torch.no_grad():
            model = Foo().eval()
            optimized_model = optimization.optimize_for_inference(model)
            torch.testing.assert_close(model(inp), optimized_model(inp))

            optimized_model2 = optimization.optimize_for_inference(
                model, pass_config={"remove_dropout": False}
            )
            torch.testing.assert_close(model(inp), optimized_model2(inp))
//  [python]
def test_padded_tensor(use_nnutils):
    x = torch.tensor(
        [
            [
                [
                    [1, 2, 3, 4, 5, 6, 7, 8],
                    [9, 10, 11, 12, 13, 14, 15, 16],
                    [17, 18, 19, 20, 21, 22, 23, 24],
                    [25, 26, 27, 28, 29, 30, 31, 32],
                ]
            ]
        ],
        dtype=torch.double,
        requires_grad=True,
    )
    xs = torch.tensor([[3, 4]])
    layer = PyramidMaxPool2d(levels=[1, 2], use_nnutils=use_nnutils)
    y = layer(PaddedTensor(x, xs))
    torch.testing.assert_allclose(
        y, torch.tensor([[20, 10, 12, 18, 20]], dtype=x.dtype)
    )
    torch.autograd.gradcheck(lambda x: layer(PaddedTensor(x, xs)), x)
//  [python]
def test_padded_tensor(use_nnutils):
    x = torch.tensor(
        [
            [
                [
                    [1, 2, 3, 4, 5, 6, 7, 8],
                    [9, 10, 11, 12, 13, 14, 15, 16],
                    [17, 18, 19, 20, 21, 22, 23, 24],
                    [25, 26, 27, 28, 29, 30, 31, 32],
                ]
            ]
        ],
        dtype=torch.double,
        requires_grad=True,
    )
    xs = torch.tensor([[3, 4]])
    layer = PyramidMaxPool2d(levels=[1, 2], use_nnutils=use_nnutils)
    y = layer(PaddedTensor(x, xs))
    torch.testing.assert_close(
        y, torch.tensor([[20, 10, 12, 18, 20]], dtype=x.dtype)
    )
    torch.autograd.gradcheck(lambda x: layer(PaddedTensor(x, xs)), x)
//  [python]
def test_normalize_binary_operators(self):
        ops_to_test = {
            torch.add,
            torch.mul,
            torch.sub,
            torch.div,
            torch.floor_divide,
            torch.remainder,
            torch.eq,
            torch.ne,
            torch.lt,
            torch.le,
            torch.gt,
            torch.ge,
        }

        # Test Tensor/Tensor callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x, y):
                    return op(x, y)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x, y = torch.randn(3, 4), torch.randn(3, 4)
            torch.testing.assert_allclose(traced(x, y), normalized(x, y))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )

        # Test Tensor/scalar callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x):
                    return op(x, 42)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x = torch.randn(3, 4)
            torch.testing.assert_allclose(traced(x), normalized(x))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )
//  [python]
def test_normalize_binary_operators(self):
        ops_to_test = {
            torch.add,
            torch.mul,
            torch.sub,
            torch.div,
            torch.floor_divide,
            torch.remainder,
            torch.eq,
            torch.ne,
            torch.lt,
            torch.le,
            torch.gt,
            torch.ge,
        }

        # Test Tensor/Tensor callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x, y):
                    return op(x, y)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x, y = torch.randn(3, 4), torch.randn(3, 4)
            torch.testing.assert_close(traced(x, y), normalized(x, y))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )

        # Test Tensor/scalar callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x):
                    return op(x, 42)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x = torch.randn(3, 4)
            torch.testing.assert_close(traced(x), normalized(x))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )
//  [python]
def test_forward(self):
        x = torch.tensor(
            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=torch.float
        )
        m = ImageToSequence(columnwise=True)
        y = m(x)
        expected_y = torch.tensor(
            [[[1, 4, 7, 10]], [[2, 5, 8, 11]], [[3, 6, 9, 12]]], dtype=torch.float
        )
        torch.testing.assert_allclose(y, expected_y)
//  [python]
def test_forward(self):
        x = torch.tensor(
            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=torch.float
        )
        m = ImageToSequence(columnwise=True)
        y = m(x)
        expected_y = torch.tensor(
            [[[1, 4, 7, 10]], [[2, 5, 8, 11]], [[3, 6, 9, 12]]], dtype=torch.float
        )
        torch.testing.assert_close(y, expected_y)
//  [python]
def test_backward(self):
        x = torch.tensor(
            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
            dtype=torch.float,
            requires_grad=True,
        )
        m = ImageToSequence(columnwise=True)
        y = m(x)
        (dx,) = torch.autograd.grad([y.sum()], [x])
        torch.testing.assert_allclose(dx, torch.ones(4, 3))
//  [python]
def test_backward(self):
        x = torch.tensor(
            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
            dtype=torch.float,
            requires_grad=True,
        )
        m = ImageToSequence(columnwise=True)
        y = m(x)
        (dx,) = torch.autograd.grad([y.sum()], [x])
        torch.testing.assert_close(dx, torch.ones(4, 3))
//  [python]
def test_elbo():
    elbo = ELBO(2, constrain=False).double()
    elbo.μ[:] = torch.tensor([1., 2.])
    elbo.Σ[:] = torch.tensor([[1., .5], [.5, 2.]])

    approx = torch.distributions.MultivariateNormal(elbo.μ, elbo.Σ)
    s = approx.sample((100000,))

    # Test entropy
    expected = -approx.log_prob(s).mean()
    torch.testing.assert_allclose(expected, elbo.entropy(), rtol=.01, atol=.01)

    # Test prior
    prior = torch.distributions.MultivariateNormal(torch.full((2,), μ0), σ0**2 * torch.eye(2))
    expected = prior.log_prob(s).mean()
    torch.testing.assert_allclose(expected, elbo.expected_prior().sum(), rtol=.01, atol=.01)

    # Test likelihood
    n = torch.tensor([[0, 3], [3, 0]])
    w = torch.tensor([[0, 1], [1, 0]])

    s = torch.distributions.MultivariateNormal(elbo.μ, elbo.Σ).sample((100000,))
    d = s[:, :, None] - s[:, None, :]
    r = 1/(1 + torch.exp(-d))
    log_likelihood = torch.distributions.Binomial(n, r).log_prob(w.float())

    expected = log_likelihood.mean(0)

    torch.testing.assert_allclose(expected, elbo.expected_log_likelihood(n, w), rtol=.01, atol=.01)
//  [python]
def test_elbo():
    elbo = ELBO(2, constrain=False).double()
    elbo.μ[:] = torch.tensor([1., 2.])
    elbo.Σ[:] = torch.tensor([[1., .5], [.5, 2.]])

    approx = torch.distributions.MultivariateNormal(elbo.μ, elbo.Σ)
    s = approx.sample((100000,))

    # Test entropy
    expected = -approx.log_prob(s).mean()
    torch.testing.assert_close(expected, elbo.entropy(), rtol=.01, atol=.01)

    # Test prior
    prior = torch.distributions.MultivariateNormal(torch.full((2,), μ0), σ0**2 * torch.eye(2))
    expected = prior.log_prob(s).mean()
    torch.testing.assert_close(expected, elbo.expected_prior().sum(), rtol=.01, atol=.01)

    # Test likelihood
    n = torch.tensor([[0, 3], [3, 0]])
    w = torch.tensor([[0, 1], [1, 0]])

    s = torch.distributions.MultivariateNormal(elbo.μ, elbo.Σ).sample((100000,))
    d = s[:, :, None] - s[:, None, :]
    r = 1/(1 + torch.exp(-d))
    log_likelihood = torch.distributions.Binomial(n, r).log_prob(w.float())

    expected = log_likelihood.mean(0)

    torch.testing.assert_close(expected, elbo.expected_log_likelihood(n, w), rtol=.01, atol=.01)
//  [python]
def test_forward_with_size(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]], dtype=torch.float
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True)
        y, ys = m(PaddedTensor(x, xs))
        expected_y = torch.tensor(
            [[[1, 4], [7, 10]], [[2, 5], [8, 11]], [[3, 6], [0, 0]]], dtype=torch.float
        )
        torch.testing.assert_allclose(y, expected_y)
        self.assertEqual(ys, [3, 2])
//  [python]
def test_forward_with_size(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]], dtype=torch.float
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True)
        y, ys = m(PaddedTensor(x, xs))
        expected_y = torch.tensor(
            [[[1, 4], [7, 10]], [[2, 5], [8, 11]], [[3, 6], [0, 0]]], dtype=torch.float
        )
        torch.testing.assert_close(y, expected_y)
        self.assertEqual(ys, [3, 2])
//  [python]
def test_modulated_deform_conv():
    check_backend(Backend.TORCHSCRIPT)
    from mmdeploy.backend.torchscript import ops_available

    if not ops_available():
        pytest.skip('torchscript custom ops is required.')

    from mmcv.ops import ModulatedDeformConv2dPack

    from mmdeploy.apis.torch_jit import trace

    model = ModulatedDeformConv2dPack(3, 1, 1).eval()
    x = torch.rand(1, 3, 16, 16)

    jit_model = trace(model, x, None, backend='torchscript')

    out = model(x)
    jit_out = jit_model(x)

    torch.testing.assert_allclose(out, jit_out)
//  [python]
def test_modulated_deform_conv():
    check_backend(Backend.TORCHSCRIPT)
    from mmdeploy.backend.torchscript import ops_available

    if not ops_available():
        pytest.skip('torchscript custom ops is required.')

    from mmcv.ops import ModulatedDeformConv2dPack

    from mmdeploy.apis.torch_jit import trace

    model = ModulatedDeformConv2dPack(3, 1, 1).eval()
    x = torch.rand(1, 3, 16, 16)

    jit_model = trace(model, x, None, backend='torchscript')

    out = model(x)
    jit_out = jit_model(x)

    torch.testing.assert_close(out, jit_out)
//  [python]
def test_backward_with_size(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]],
            dtype=torch.float,
            requires_grad=True,
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True)
        y, ys = m(PaddedTensor(x, xs))
        (dx,) = torch.autograd.grad(
            [y[0, :, :].sum() + y[1, :, :].sum() + y[2, 0, :].sum()], [x]
        )
        expected_dx = torch.tensor(
            [[[[1, 1, 1], [1, 1, 1]]], [[[1, 1, 0], [1, 1, 0]]]], dtype=torch.float
        )
        torch.testing.assert_allclose(dx, expected_dx)
//  [python]
def test_backward_with_size(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]],
            dtype=torch.float,
            requires_grad=True,
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True)
        y, ys = m(PaddedTensor(x, xs))
        (dx,) = torch.autograd.grad(
            [y[0, :, :].sum() + y[1, :, :].sum() + y[2, 0, :].sum()], [x]
        )
        expected_dx = torch.tensor(
            [[[[1, 1, 1], [1, 1, 1]]], [[[1, 1, 0], [1, 1, 0]]]], dtype=torch.float
        )
        torch.testing.assert_close(dx, expected_dx)
//  [python]
def test_reparametrization_trick(mean, cov):
    for seed in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
        # Sampling the the PyTorch distribution class
        distr_mvn = MultivariateNormal(mean, cov)
        to.manual_seed(seed)
        smpl_distr = distr_mvn.sample()

        # The reparametrization trick done by PyTorch
        to.manual_seed(seed)
        smpl_distr_reparam = distr_mvn.sample()

        # The reparametrization trick done by hand
        to.manual_seed(seed)
        smpl_reparam = mean + to.cholesky(cov, upper=False).mv(to.randn_like(mean))

        to.testing.assert_allclose(smpl_distr, smpl_distr_reparam)
        to.testing.assert_allclose(smpl_distr, smpl_reparam)
        to.testing.assert_allclose(smpl_distr_reparam, smpl_reparam)
//  [python]
def test_reparametrization_trick(mean, cov):
    for seed in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
        # Sampling the the PyTorch distribution class
        distr_mvn = MultivariateNormal(mean, cov)
        to.manual_seed(seed)
        smpl_distr = distr_mvn.sample()

        # The reparametrization trick done by PyTorch
        to.manual_seed(seed)
        smpl_distr_reparam = distr_mvn.sample()

        # The reparametrization trick done by hand
        to.manual_seed(seed)
        smpl_reparam = mean + to.cholesky(cov, upper=False).mv(to.randn_like(mean))

        to.testing.assert_close(smpl_distr, smpl_distr_reparam)
        to.testing.assert_close(smpl_distr, smpl_reparam)
        to.testing.assert_close(smpl_distr_reparam, smpl_reparam)
//  [python]
def test_burn_single_source():
    t1 = torch.tensor([1., 2, 4, 7, 1, 3, 7, 8, 0, 3, 4])
    t2 = torch.clone(t1)
    t2[:4] += 2

    sdr = ci_sdr.pt.ci_sdr(t1, t2, filter_length=3)
    assert sdr.shape == (), sdr.shape
    torch.testing.assert_allclose(sdr, 13.592828750610352)
//  [python]
def test_burn_single_source():
    t1 = torch.tensor([1., 2, 4, 7, 1, 3, 7, 8, 0, 3, 4])
    t2 = torch.clone(t1)
    t2[:4] += 2

    sdr = ci_sdr.pt.ci_sdr(t1, t2, filter_length=3)
    assert sdr.shape == (), sdr.shape
    torch.testing.assert_close(sdr, 13.592828750610352)
//  [python]
def test_forward_backward_packed(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]],
            dtype=torch.float,
            requires_grad=True,
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True, return_packed=True)
        # Test forward
        y = m(PaddedTensor(x, xs))
        expected_y = torch.tensor(
            [[1, 4], [7, 10], [2, 5], [8, 11], [3, 6]], dtype=torch.float
        )
        torch.testing.assert_allclose(y.data, expected_y)
        self.assertTrue(torch.equal(torch.tensor([2, 2, 1]), y.batch_sizes))
        # Test backward pass
        (dx,) = torch.autograd.grad([y.data.sum()], [x])
        expected_dx = torch.tensor(
            [[[[1, 1, 1], [1, 1, 1]]], [[[1, 1, 0], [1, 1, 0]]]], dtype=torch.float
        )
        torch.testing.assert_allclose(dx, expected_dx)
//  [python]
def test_forward_backward_packed(self):
        x = torch.tensor(
            [[[[1, 2, 3], [4, 5, 6]]], [[[7, 8, 0], [10, 11, 0]]]],
            dtype=torch.float,
            requires_grad=True,
        )
        xs = torch.tensor([[2, 3], [2, 2]])
        m = ImageToSequence(columnwise=True, return_packed=True)
        # Test forward
        y = m(PaddedTensor(x, xs))
        expected_y = torch.tensor(
            [[1, 4], [7, 10], [2, 5], [8, 11], [3, 6]], dtype=torch.float
        )
        torch.testing.assert_close(y.data, expected_y)
        self.assertTrue(torch.equal(torch.tensor([2, 2, 1]), y.batch_sizes))
        # Test backward pass
        (dx,) = torch.autograd.grad([y.data.sum()], [x])
        expected_dx = torch.tensor(
            [[[[1, 1, 1], [1, 1, 1]]], [[[1, 1, 0], [1, 1, 0]]]], dtype=torch.float
        )
        torch.testing.assert_close(dx, expected_dx)
//  [python]
def test_burn_multi_source():
    t1 = torch.tensor([
        [1., 2, 4, 7, 1, 3, 7, 8, 0, 3, 4],
        [5., 2, 7, 9, 3, 8, 4, 2, 9, 4, 5],
    ])
    t2 = torch.clone(t1)
    t2[:, :4] += 2

    sdr = ci_sdr.pt.ci_sdr(t1, t2, filter_length=3, compute_permutation=False)
    assert sdr.shape == (2,), sdr.shape
    torch.testing.assert_allclose(sdr, [13.592828750610352, 17.48115348815918])

    sdr = ci_sdr.pt.ci_sdr(
        t1, t2[(1, 0), :], filter_length=3, compute_permutation=True)
    assert sdr.shape == (2,), sdr.shape
    torch.testing.assert_allclose(sdr, [13.592828750610352, 17.48115348815918])
//  [python]
def test_burn_multi_source():
    t1 = torch.tensor([
        [1., 2, 4, 7, 1, 3, 7, 8, 0, 3, 4],
        [5., 2, 7, 9, 3, 8, 4, 2, 9, 4, 5],
    ])
    t2 = torch.clone(t1)
    t2[:, :4] += 2

    sdr = ci_sdr.pt.ci_sdr(t1, t2, filter_length=3, compute_permutation=False)
    assert sdr.shape == (2,), sdr.shape
    torch.testing.assert_close(sdr, [13.592828750610352, 17.48115348815918])

    sdr = ci_sdr.pt.ci_sdr(
        t1, t2[(1, 0), :], filter_length=3, compute_permutation=True)
    assert sdr.shape == (2,), sdr.shape
    torch.testing.assert_close(sdr, [13.592828750610352, 17.48115348815918])
//  [python]
def test_vtrace(self):
        """Tests V-trace against ground truth data calculated in python."""
        batch_size = 5
        seq_len = 5

        # Create log_rhos such that rho will span from near-zero to above the
        # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5),
        # so that rho is in approx [0.08, 12.2).
        log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len)
        log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
        values = {
            'behavior_action_log_probs': torch.zeros_like(log_rhos),
            'target_action_log_probs': log_rhos,
            # T, B where B_i: [0.9 / (i+1)] * T
            'discounts': torch.tensor(
                np.array([[0.9 / (b + 1) for b in range(batch_size)] for _ in range(seq_len)]), dtype=torch.float32
            ),
            'rewards': _shaped_arange(seq_len, batch_size),
            'values': _shaped_arange(seq_len, batch_size) / batch_size,
            'bootstrap_value': _shaped_arange(batch_size) + 1.0,
            'clip_rho_threshold': 3.7,
            'clip_pg_rho_threshold': 2.2,
        }

        output = vtrace.from_importance_weights(**values)
        ground_truth_v = _ground_truth_calculation(**values)

        torch.testing.assert_allclose(output.vs, ground_truth_v.vs)
        torch.testing.assert_allclose(output.pg_advantages, ground_truth_v.pg_advantages)
//  [python]
def test_vtrace(self):
        """Tests V-trace against ground truth data calculated in python."""
        batch_size = 5
        seq_len = 5

        # Create log_rhos such that rho will span from near-zero to above the
        # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5),
        # so that rho is in approx [0.08, 12.2).
        log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len)
        log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
        values = {
            'behavior_action_log_probs': torch.zeros_like(log_rhos),
            'target_action_log_probs': log_rhos,
            # T, B where B_i: [0.9 / (i+1)] * T
            'discounts': torch.tensor(
                np.array([[0.9 / (b + 1) for b in range(batch_size)] for _ in range(seq_len)]), dtype=torch.float32
            ),
            'rewards': _shaped_arange(seq_len, batch_size),
            'values': _shaped_arange(seq_len, batch_size) / batch_size,
            'bootstrap_value': _shaped_arange(batch_size) + 1.0,
            'clip_rho_threshold': 3.7,
            'clip_pg_rho_threshold': 2.2,
        }

        output = vtrace.from_importance_weights(**values)
        ground_truth_v = _ground_truth_calculation(**values)

        torch.testing.assert_close(output.vs, ground_truth_v.vs)
        torch.testing.assert_close(output.pg_advantages, ground_truth_v.pg_advantages)
//  [python]
def test_deit_transformers():

    classes = [
        (toroidal.models.DeiTBase, timm.models.deit_base_patch16_224),
        (toroidal.models.DeiTTiny, timm.models.deit_tiny_patch16_224),
    ]
    for cls, cls_ref in classes:
        model_ref = cls_ref(pretrained=True)

        model = cls()

        sd = model_ref.state_dict()
        sd["patch_embedding.class_token"] = sd.pop("cls_token")
        sd["patch_embedding.pos_embedding"] = sd.pop("pos_embed")
        sd["patch_embedding.patch_to_vec.weight"] = sd.pop("patch_embed.proj.weight")
        sd["patch_embedding.patch_to_vec.bias"] = sd.pop("patch_embed.proj.bias")

        for k in list(sd.keys()):
            if ".mlp.fc1" in k:
                sd[k.replace(".mlp.fc1", ".mlp.0")] = sd.pop(k)
            elif ".mlp.fc2" in k:
                sd[k.replace(".mlp.fc2", ".mlp.2")] = sd.pop(k)

        model.load_state_dict(sd)

        for training in (True, False):
            model.train(training)
            model_ref.train(training)

            inp = torch.randn(1, 3, 224, 224)
            torch.manual_seed(1234)
            ref_out = model_ref(inp)
            torch.manual_seed(1234)
            out = model(inp)

            torch.testing.assert_allclose(ref_out, out)
//  [python]
def test_deit_transformers():

    classes = [
        (toroidal.models.DeiTBase, timm.models.deit_base_patch16_224),
        (toroidal.models.DeiTTiny, timm.models.deit_tiny_patch16_224),
    ]
    for cls, cls_ref in classes:
        model_ref = cls_ref(pretrained=True)

        model = cls()

        sd = model_ref.state_dict()
        sd["patch_embedding.class_token"] = sd.pop("cls_token")
        sd["patch_embedding.pos_embedding"] = sd.pop("pos_embed")
        sd["patch_embedding.patch_to_vec.weight"] = sd.pop("patch_embed.proj.weight")
        sd["patch_embedding.patch_to_vec.bias"] = sd.pop("patch_embed.proj.bias")

        for k in list(sd.keys()):
            if ".mlp.fc1" in k:
                sd[k.replace(".mlp.fc1", ".mlp.0")] = sd.pop(k)
            elif ".mlp.fc2" in k:
                sd[k.replace(".mlp.fc2", ".mlp.2")] = sd.pop(k)

        model.load_state_dict(sd)

        for training in (True, False):
            model.train(training)
            model_ref.train(training)

            inp = torch.randn(1, 3, 224, 224)
            torch.manual_seed(1234)
            ref_out = model_ref(inp)
            torch.manual_seed(1234)
            out = model(inp)

            torch.testing.assert_close(ref_out, out)
//  [python]
def test_adjoint(A, AT):
    x = torch.randn(10, 10, 10, 10)
    y = torch.randn(10, 10, 10, 10)

    torch.testing.assert_allclose(
        torch.dot(A(x).ravel(), y.ravel()),
        torch.dot(x.ravel(), AT(y).ravel())
    )
//  [python]
def test_adjoint(A, AT):
    x = torch.randn(10, 10, 10, 10)
    y = torch.randn(10, 10, 10, 10)

    torch.testing.assert_close(
        torch.dot(A(x).ravel(), y.ravel()),
        torch.dot(x.ravel(), AT(y).ravel())
    )
//  [python]
def test_standard_conv_cpu():
    '''Simply compares our result to PyTorch's implementation'''
    input = torch.ones(params.bs, params.in_channels, 8, 8).double()
    input.requires_grad = True

    pytorch_layer = torch.nn.Conv2d(in_channels,
                                    out_channels,
                                    kernel_size=3,
                                    stride=2,
                                    padding=1,
                                    dilation=2).double()
    my_layer = Convolution(in_channels,
                           out_channels,
                           kernel_size=3,
                           stride=2,
                           padding=1,
                           dilation=2).double()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    testing.assert_allclose(my_output, pytorch_output)
//  [python]
def test_standard_conv_cpu():
    '''Simply compares our result to PyTorch's implementation'''
    input = torch.ones(params.bs, params.in_channels, 8, 8).double()
    input.requires_grad = True

    pytorch_layer = torch.nn.Conv2d(in_channels,
                                    out_channels,
                                    kernel_size=3,
                                    stride=2,
                                    padding=1,
                                    dilation=2).double()
    my_layer = Convolution(in_channels,
                           out_channels,
                           kernel_size=3,
                           stride=2,
                           padding=1,
                           dilation=2).double()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    torch.testing.assert_close(my_output, pytorch_output)
//  [python]
def test_standard_conv_cuda():
    '''Simply compares our result to PyTorch's implementation'''

    input = torch.ones(params.bs, params.in_channels, 8, 8).double().cuda()
    input.requires_grad = True

    pytorch_layer = torch.nn.Conv2d(in_channels,
                                    out_channels,
                                    kernel_size=3,
                                    stride=2,
                                    padding=1,
                                    dilation=2).double().cuda()
    my_layer = Convolution(in_channels,
                           out_channels,
                           kernel_size=3,
                           stride=2,
                           padding=1,
                           dilation=2).double().cuda()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    testing.assert_allclose(my_output, pytorch_output)
//  [python]
def test_standard_conv_cuda():
    '''Simply compares our result to PyTorch's implementation'''

    input = torch.ones(params.bs, params.in_channels, 8, 8).double().cuda()
    input.requires_grad = True

    pytorch_layer = torch.nn.Conv2d(in_channels,
                                    out_channels,
                                    kernel_size=3,
                                    stride=2,
                                    padding=1,
                                    dilation=2).double().cuda()
    my_layer = Convolution(in_channels,
                           out_channels,
                           kernel_size=3,
                           stride=2,
                           padding=1,
                           dilation=2).double().cuda()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    torch.testing.assert_close(my_output, pytorch_output)
//  [python]
def test_invert():
    # Define a simple bijector to invert
    ar = Compose(
        [
            AffineAutoregressive(params.DenseAutoregressive()),
            AffineAutoregressive(params.DenseAutoregressive()),
        ]
    )
    shape = torch.Size(
        [
            16,
        ]
    )

    # Instantiate the bijector and its inverse
    bij = ar(shape=shape)
    inv_bij = bijectors.Invert(ar)(shape=shape)

    # Make parameters the same for both
    inv_bij.load_state_dict(bij.state_dict(prefix="bijector."))

    # Test if inversion is correct
    x = torch.randn(50, 16, requires_grad=True)
    torch.testing.assert_allclose(inv_bij.forward(x), bij.inverse(x))

    y = inv_bij.forward(x)

    # checks that no warning is displayed, which can happen if no cache is used
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        inv_bij.log_abs_det_jacobian(x, y)

    with pytest.warns(UserWarning):
        y_det = y.detach_from_flow()
        inv_bij.log_abs_det_jacobian(x, y_det)

    y = y.detach_from_flow()
    torch.testing.assert_allclose(
        inv_bij.log_abs_det_jacobian(x, y), bij.log_abs_det_jacobian(y, x)
    )
//  [python]
def test_invert():
    # Define a simple bijector to invert
    ar = Compose(
        [
            AffineAutoregressive(params.DenseAutoregressive()),
            AffineAutoregressive(params.DenseAutoregressive()),
        ]
    )
    shape = torch.Size(
        [
            16,
        ]
    )

    # Instantiate the bijector and its inverse
    bij = ar(shape=shape)
    inv_bij = bijectors.Invert(ar)(shape=shape)

    # Make parameters the same for both
    inv_bij.load_state_dict(bij.state_dict(prefix="bijector."))

    # Test if inversion is correct
    x = torch.randn(50, 16, requires_grad=True)
    torch.testing.assert_close(inv_bij.forward(x), bij.inverse(x))

    y = inv_bij.forward(x)

    # checks that no warning is displayed, which can happen if no cache is used
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        inv_bij.log_abs_det_jacobian(x, y)

    with pytest.warns(UserWarning):
        y_det = y.detach_from_flow()
        inv_bij.log_abs_det_jacobian(x, y_det)

    y = y.detach_from_flow()
    torch.testing.assert_close(
        inv_bij.log_abs_det_jacobian(x, y), bij.log_abs_det_jacobian(y, x)
    )
//  [python]
def test_standard_transposed_conv_cpu():
    '''Simply compares our result to PyTorch's implementation'''
    input = torch.ones(params.bs, params.in_channels, 5, 5).double()
    input.requires_grad = True

    pytorch_layer = torch.nn.ConvTranspose2d(params.in_channels,
                                             params.out_channels,
                                             kernel_size=3,
                                             stride=2,
                                             padding=1,
                                             dilation=2).double()
    my_layer = TransposedConvolution(params.in_channels,
                                     params.out_channels,
                                     kernel_size=3,
                                     stride=2,
                                     padding=1,
                                     dilation=2).double()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    testing.assert_allclose(my_output, pytorch_output)