jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_encoder_batch(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
        fake_sample = torch.rand(5, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_sample[2, 8:, :] = 0
        fake_sample[3, 8:, :] = 0
        fake_sample[4, 5:, :] = 0
        fake_lengths = torch.LongTensor([27, 13, 8, 8, 5])
        encoder = ConformerEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)

        def test_item(item_idx):
            item_len = fake_lengths[item_idx].item()
            item_out_len = math.ceil(item_len / 4)
            fake_sample2 = fake_sample[item_idx, :item_len, :]
            net_out2 = encoder.forward(
                fake_sample2.unsqueeze(0), fake_lengths[item_idx].unsqueeze(0), return_all_hiddens=True)
            torch.testing.assert_allclose(
                    net_out["encoder_out"][0][:item_out_len, item_idx, :],
                    net_out2["encoder_out"][0][:, 0, :])

        for i in range(5):
            test_item(i)
//  [python]
def test_encoder_batch(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
        fake_sample = torch.rand(5, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_sample[2, 8:, :] = 0
        fake_sample[3, 8:, :] = 0
        fake_sample[4, 5:, :] = 0
        fake_lengths = torch.LongTensor([27, 13, 8, 8, 5])
        encoder = ConformerEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)

        def test_item(item_idx):
            item_len = fake_lengths[item_idx].item()
            item_out_len = math.ceil(item_len / 4)
            fake_sample2 = fake_sample[item_idx, :item_len, :]
            net_out2 = encoder.forward(
                fake_sample2.unsqueeze(0), fake_lengths[item_idx].unsqueeze(0), return_all_hiddens=True)
            torch.testing.assert_close(
                    net_out["encoder_out"][0][:item_out_len, item_idx, :],
                    net_out2["encoder_out"][0][:, 0, :])

        for i in range(5):
            test_item(i)
//  [python]
def test_unresample_nearest_out_of_bounds_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_map8(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4, 0]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    testing.assert_allclose(output, correct_output)
//  [python]
def test_unresample_nearest_out_of_bounds_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_map8(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4, 0]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_bilinear_integer_sampling_cpu():

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

    # 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_bilinear_real_sampling_cpu():

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

    # 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([[0.5, 5.5], [9.5, 5.5], [1.5,
                                                            10.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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([[0.5, 5.5], [9.5, 5.5], [1.5,
                                                            10.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=False)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=False)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_bispherical_integer_sampling_cpu():

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

    # 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_bispherical_real_sampling_cpu():

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

    # 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([[0.5, 5.5], [9.5, 9.5], [1.5,
                                                            10.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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([[0.5, 5.5], [9.5, 9.5], [1.5,
                                                            10.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=False)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=False)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

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

    # Manually computed correct result
    correct_output = torch.tensor([[14, 15.8], [9.6, 19.1]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

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

    # Manually computed correct result
    correct_output = torch.tensor([[14, 15.8], [9.6, 19.1]]).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_cuda():

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0.5, 5.5], [9.5, 5.5],
                                   [1.5, 10.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0.5, 5.5], [9.5, 5.5],
                                   [1.5, 10.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [3, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [3, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

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

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4.75, 11]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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_map8(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0, 0], [0, 0], [4.75, 11]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0.5, 5.5], [9.5, 9.5],
                                   [1.5, 10.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # 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=True)

    # Manually computed correct result
    correct_output = torch.tensor([[0.5, 5.5], [9.5, 9.5],
                                   [1.5, 10.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_4x7().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map5(),
        interp_weights=params.interp_weights0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 15.8], [9.6, 19.1]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

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

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_resample_test(
        layer,
        input=params.input_4x7().repeat(bs, channels, 1, 1),
        sample_map=params.sample_map5(),
        interp_weights=params.interp_weights0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 15.8], [9.6, 19.1]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MaxPool layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[11, 19, 14, 18, 13], [18, 17, 19, 18, 19],
                                   [19, 12, 19, 10, 17], [18, 19, 10, 17,
                                                          17]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MaxPool layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[11, 19, 14, 18, 13], [18, 17, 19, 18, 19],
                                   [19, 12, 19, 10, 17], [18, 19, 10, 17,
                                                          17]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map2(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[19, 17], [19, 19]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map2(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[19, 17], [19, 19]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map3(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[18.5, 16.5], [18.5, 18.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map3(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[18.5, 16.5], [18.5, 18.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map1(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 5, 8.25, 13,
                                    16], [9.25, 9, 9.25, 13, 13],
                                   [13, 15, 15, 13, 11], [13, 13, 13, 13,
                                                          15]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map1(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 5, 8.25, 13,
                                    16], [9.25, 9, 9.25, 13, 13],
                                   [13, 15, 15, 13, 11], [13, 13, 13, 13,
                                                          15]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MaxPool layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[11, 19, 14, 18, 13], [18, 17, 19, 18, 19],
                                   [19, 12, 19, 10, 17], [18, 19, 10, 17,
                                                          17]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MaxPool layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[11, 19, 14, 18, 13], [18, 17, 19, 18, 19],
                                   [19, 12, 19, 10, 17], [18, 19, 10, 17,
                                                          17]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedAvgPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor(
        [[7.5, 6.25, 7.75, 9.75, 8.25], [12.25, 10., 10., 13.5, 10.75],
         [11.5, 8.75, 11.75, 6.5, 8.25], [12.5, 9., 6.75, 10., 11.25]
         ], ).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedAvgPool(kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=False)

    # Manually computed correct result
    correct_output = torch.tensor(
        [[7.5, 6.25, 7.75, 9.75, 8.25], [12.25, 10., 10., 13.5, 10.75],
         [11.5, 8.75, 11.75, 6.5, 8.25], [12.5, 9., 6.75, 10., 11.25]
         ], ).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map1(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 5, 8.25, 13,
                                    16], [9.25, 9, 9.25, 13, 13],
                                   [13, 15, 15, 13, 11], [13, 13, 13, 13,
                                                          15]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map1(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[14, 5, 8.25, 13,
                                    16], [9.25, 9, 9.25, 13, 13],
                                   [13, 15, 15, 13, 11], [13, 13, 13, 13,
                                                          15]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    gpu_mem_0 = report_memory("Before model init")
    with slapo.init_empty_weights(enable=True):
        _ = Model()
    gpu_mem_1 = report_memory("After model init")
    if ngpu == "multi":
        assert (
            gpu_mem_1 == gpu_mem_0
        ), f"GPU memory 0: {gpu_mem_0}, GPU memory 1: {gpu_mem_1}"

    model = Model()
    sch = slapo.create_schedule(copy.deepcopy(model))
    if ngpu == "multi":
        # Tensor parallelism.
        sch["linear1"].shard("weight", axis=0)
        sch["linear1"].shard("bias", axis=0)
        sch["linear1"].sync(mode="bwd_post", sync_op_or_fn="all_reduce")
        sch["linear2"].shard("weight", axis=1)
        sch["linear2"].sync(mode="fwd_post", sync_op_or_fn="all_reduce")
    else:
        # Data parallelism.
        pass
    sch_model, _ = slapo.build(sch, init_weights=False)

    model.cuda(local_rank)
    sch_model.cuda(local_rank)
    torch.testing.assert_allclose(
        sch_model.linear1.weight,
        get_partial_tensor(model.linear1.weight, 0, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear1.weight,
    )
    torch.testing.assert_allclose(
        sch_model.linear1.bias,
        get_partial_tensor(model.linear1.bias, 0, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear1.bias,
    )
    torch.testing.assert_allclose(
        sch_model.linear2.weight,
        get_partial_tensor(model.linear2.weight, 1, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear2.weight,
    )
//  [python]
def test_consolidation_default_init(init_dist, ngpu):

    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    gpu_mem_0 = report_memory("Before model init")
    with slapo.init_empty_weights(enable=True):
        _ = Model()
    gpu_mem_1 = report_memory("After model init")
    if ngpu == "multi":
        assert (
            gpu_mem_1 == gpu_mem_0
        ), f"GPU memory 0: {gpu_mem_0}, GPU memory 1: {gpu_mem_1}"

    model = Model()
    sch = slapo.create_schedule(copy.deepcopy(model))
    if ngpu == "multi":
        # Tensor parallelism.
        sch["linear1"].shard("weight", axis=0)
        sch["linear1"].shard("bias", axis=0)
        sch["linear1"].sync(mode="bwd_post", sync_op_or_fn="all_reduce")
        sch["linear2"].shard("weight", axis=1)
        sch["linear2"].sync(mode="fwd_post", sync_op_or_fn="all_reduce")
    else:
        # Data parallelism.
        pass
    sch_model, _ = slapo.build(sch, init_weights=False)

    model.cuda(local_rank)
    sch_model.cuda(local_rank)
    torch.testing.assert_close(
        sch_model.linear1.weight,
        get_partial_tensor(model.linear1.weight, 0, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear1.weight,
    )
    torch.testing.assert_close(
        sch_model.linear1.bias,
        get_partial_tensor(model.linear1.bias, 0, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear1.bias,
    )
    torch.testing.assert_close(
        sch_model.linear2.weight,
        get_partial_tensor(model.linear2.weight, 1, local_rank, sch.world_size)
        if ngpu == "multi"
        else model.linear2.weight,
    )
//  [python]
def test_max_pool_integer_downsampling_cuda():

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map2(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[19, 17], [19, 19]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map2(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[19, 17], [19, 19]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map3(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[18.5, 16.5], [18.5, 18.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedMaxPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map3(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor([[18.5, 16.5], [18.5, 18.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_kwinners_grad(self):
        n = 3
        x = torch.tensor([[0, 2, 1],
                          [14, 13, 12]], dtype=torch.float, requires_grad=True)
        grad = torch.tensor([[5, 6, 7],
                             [45, 46, 47]], dtype=torch.float)
        expected = torch.tensor([[0, 6, 0],
                                 [45, 0, 0]], dtype=torch.float)

        for break_ties in [True, False]:
            with self.subTest(break_ties=break_ties):
                kw = KWinners(n,
                              percent_on=(1 / 3),
                              k_inference_factor=1.0,
                              boost_strength=0.0,
                              break_ties=break_ties)
                kw(x).backward(grad)
                torch.testing.assert_allclose(x.grad, expected)
                x.grad.zero_()
//  [python]
def test_kwinners_grad(self):
        n = 3
        x = torch.tensor([[0, 2, 1],
                          [14, 13, 12]], dtype=torch.float, requires_grad=True)
        grad = torch.tensor([[5, 6, 7],
                             [45, 46, 47]], dtype=torch.float)
        expected = torch.tensor([[0, 6, 0],
                                 [45, 0, 0]], dtype=torch.float)

        for break_ties in [True, False]:
            with self.subTest(break_ties=break_ties):
                kw = KWinners(n,
                              percent_on=(1 / 3),
                              k_inference_factor=1.0,
                              boost_strength=0.0,
                              break_ties=break_ties)
                kw(x).backward(grad)
                torch.testing.assert_close(x.grad, expected)
                x.grad.zero_()
//  [python]
def test_maskedconvk3(self):
        torch.testing.assert_allclose(self.convk3(self.data),
                                      torch.tensor([[5, -6, 12, -10, 24],
                                                    [-6, 30, -13, 35, 1],
                                                    [-5, 8, -14, 19, -8],
                                                    [31, -13, 35, -3, -7],
                                                    [13, -8, 20, -12, 32]],
                                                   dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_maskedconvk3(self):
        torch.testing.assert_close(self.convk3(self.data),
                                      torch.tensor([[5, -6, 12, -10, 24],
                                                    [-6, 30, -13, 35, 1],
                                                    [-5, 8, -14, 19, -8],
                                                    [31, -13, 35, -3, -7],
                                                    [13, -8, 20, -12, 32]],
                                                   dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_avg_pool_integer_sampling_cuda():

    # Basic MappedTransposedConvolution layer
    layer = MappedAvgPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor(
        [[7.5, 6.25, 7.75, 9.75, 8.25], [12.25, 10., 10., 13.5, 10.75],
         [11.5, 8.75, 11.75, 6.5, 8.25], [12.5, 9., 6.75, 10., 11.25]
         ], ).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedAvgPool(kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_pool_test(
        layer,
        input=params.input_4x5().repeat(bs, 1, 1, 1),
        sample_map=params.sample_map0(),
        cuda=True)

    # Manually computed correct result
    correct_output = torch.tensor(
        [[7.5, 6.25, 7.75, 9.75, 8.25], [12.25, 10., 10., 13.5, 10.75],
         [11.5, 8.75, 11.75, 6.5, 8.25], [12.5, 9., 6.75, 10., 11.25]
         ], ).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_blur_equivalence(kernel_size, sigma, channels):
    for _ in range(10):
        input_tensor = torch.randn((3, channels, 128, 128))
        kornia = korniaGaussianBlur2d(kernel_size, sigma, separable=False)
        blur_kornia = kornia(input_tensor)
        gaussian = GaussianBlur2d(kernel_size=kernel_size, sigma=sigma, channels=channels)
        blur_gaussian = gaussian(input_tensor)
        torch.testing.assert_allclose(blur_kornia, blur_gaussian)