jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_integer_sampling_cpu():

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map2(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 0, 0, 5
    ], [2, 0, 0, 1, 0], [0, 0, 1, 1, 0], [0, 3, 1, 3, 5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map2(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 0, 0, 5
    ], [2, 0, 0, 1, 0], [0, 0, 1, 1, 0], [0, 3, 1, 3, 5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("ArthurZ/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        # fmt: off
        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])
        # fmt: on
        torch.testing.assert_allclose(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("ArthurZ/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        # fmt: off
        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])
        # fmt: on
        torch.testing.assert_close(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_deit_tiny():
    model_name = "deit_tiny_patch16_224-attn"
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
    ]
    model_config = [j for i in model_config for j in i]

    model = build_model(model_name)
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(1, 197, 192)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    torch.testing.assert_allclose(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_deit_tiny():
    model_name = "deit_tiny_patch16_224-attn"
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
    ]
    model_config = [j for i in model_config for j in i]

    model = build_model(model_name)
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(1, 197, 192)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    torch.testing.assert_close(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_bilinear_interpolation_sampling_cpu():

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map3(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor(
        [[1.0, 0.0, 0.0, 2.5, 2.5], [1.0, 0.0, 0.5, 0.5, 0.0],
         [0.0, 0.5, 1.0, 0.5, 0.0], [1.5, 2.0, 2.0, 4.0, 2.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map3(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor(
        [[1.0, 0.0, 0.0, 2.5, 2.5], [1.0, 0.0, 0.5, 0.5, 0.0],
         [0.0, 0.5, 1.0, 0.5, 0.0], [1.5, 2.0, 2.0, 4.0, 2.5]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x3().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map25(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 3, 0, 13
    ], [8, 0, 5, 6, 0], [8, 0, 8, 0, 4], [0, 0, 3, 0, 0]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(in_channels=in_channels,
                                        out_channels=out_channels,
                                        kernel_size=kernel_size).double()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x3().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map25(),
        cuda=False)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 3, 0, 13
    ], [8, 0, 5, 6, 0], [8, 0, 8, 0, 4], [0, 0, 3, 0, 0]]).double()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map2(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 0, 0, 5
    ], [2, 0, 0, 1, 0], [0, 0, 1, 1, 0], [0, 3, 1, 3, 5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map2(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 0, 0, 5
    ], [2, 0, 0, 1, 0], [0, 0, 1, 1, 0], [0, 3, 1, 3, 5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_uniform_random_selection(self):
        scores = torch.tensor(
            [
                [1.0, 2.0, 3.0, 4.0, 5.0],
                [5.0, 1.0, 2.0, 3.0, 4.0],
            ]
        )
        sampler = EpsilonGreedyActionSampler(epsilon=1.0)

        action = sampler.sample_action(scores)

        test_log_prob = torch.tensor(
            [-1.60944, -1.60944],
            dtype=torch.float,
        )

        torch.testing.assert_allclose(action.log_prob, test_log_prob)
//  [python]
def test_uniform_random_selection(self):
        scores = torch.tensor(
            [
                [1.0, 2.0, 3.0, 4.0, 5.0],
                [5.0, 1.0, 2.0, 3.0, 4.0],
            ]
        )
        sampler = EpsilonGreedyActionSampler(epsilon=1.0)

        action = sampler.sample_action(scores)

        test_log_prob = torch.tensor(
            [-1.60944, -1.60944],
            dtype=torch.float,
        )

        torch.testing.assert_close(action.log_prob, test_log_prob)
//  [python]
def test_greedy_selection(self):
        scores = torch.tensor(
            [
                [1.0, 2.0, 3.0, 4.0, 5.0],
                [5.0, 1.0, 2.0, 3.0, 4.0],
            ]
        )
        sampler = EpsilonGreedyActionSampler(epsilon=0.0)

        test_action = torch.tensor(
            [
                [0, 0, 0, 0, 1],
                [1, 0, 0, 0, 0],
            ],
            dtype=torch.long,
        )
        action = sampler.sample_action(scores)

        torch.testing.assert_allclose(action.action, test_action)

        test_log_prob = torch.tensor(
            [0.0, 0.0],
            dtype=torch.float,
        )

        torch.testing.assert_allclose(action.log_prob, test_log_prob)
//  [python]
def test_greedy_selection(self):
        scores = torch.tensor(
            [
                [1.0, 2.0, 3.0, 4.0, 5.0],
                [5.0, 1.0, 2.0, 3.0, 4.0],
            ]
        )
        sampler = EpsilonGreedyActionSampler(epsilon=0.0)

        test_action = torch.tensor(
            [
                [0, 0, 0, 0, 1],
                [1, 0, 0, 0, 0],
            ],
            dtype=torch.long,
        )
        action = sampler.sample_action(scores)

        torch.testing.assert_close(action.action, test_action)

        test_log_prob = torch.tensor(
            [0.0, 0.0],
            dtype=torch.float,
        )

        torch.testing.assert_close(action.log_prob, test_log_prob)
//  [python]
def test_fold(act_shape, stride_h, stride_w, device):
    torch.manual_seed(0)

    torch_input = torch.randn(act_shape, dtype=torch.bfloat16)

    expected = fold_torch(torch_input, stride_h, stride_w)
    expected = expected.reshape(1, 1, -1, expected.shape[-1])

    tt_input = torch2tt_tensor(
        torch_input,
        device,
        tt.Layout.ROW_MAJOR,
        tt_memory_config=tt.MemoryConfig(tt.TensorMemoryLayout.INTERLEAVED, tt.BufferType.L1),
    )

    tt_out = tt.fold(tt_input, stride_h, stride_w)
    actual = tt2torch_tensor(tt_out)

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_fold(act_shape, stride_h, stride_w, device):
    torch.manual_seed(0)

    torch_input = torch.randn(act_shape, dtype=torch.bfloat16)

    expected = fold_torch(torch_input, stride_h, stride_w)
    expected = expected.reshape(1, 1, -1, expected.shape[-1])

    tt_input = torch2tt_tensor(
        torch_input,
        device,
        tt.Layout.ROW_MAJOR,
        tt_memory_config=tt.MemoryConfig(tt.TensorMemoryLayout.INTERLEAVED, tt.BufferType.L1),
    )

    tt_out = tt.fold(tt_input, stride_h, stride_w)
    actual = tt2torch_tensor(tt_out)

    torch.testing.assert_close(actual, expected)
//  [python]
def test_bilinear_interpolation_sampling_cuda():

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map3(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor(
        [[1.0, 0.0, 0.0, 2.5, 2.5], [1.0, 0.0, 0.5, 0.5, 0.0],
         [0.0, 0.5, 1.0, 0.5, 0.0], [1.5, 2.0, 2.0, 4.0, 2.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x2().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map3(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor(
        [[1.0, 0.0, 0.0, 2.5, 2.5], [1.0, 0.0, 0.5, 0.5, 0.0],
         [0.0, 0.5, 1.0, 0.5, 0.0], [1.5, 2.0, 2.0, 4.0, 2.5]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x3().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map25(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 3, 0, 13
    ], [8, 0, 5, 6, 0], [8, 0, 8, 0, 4], [0, 0, 3, 0, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

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

    # Basic MappedTransposedConvolution layer
    layer = MappedTransposedConvolution(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size).double().cuda()

    # Run a forward and backward pass
    output, forward_time, backward_time, gradcheck_res = utils.mapped_transposed_conv_test(
        layer,
        weight=params.transposed_weights_unit(in_channels, out_channels),
        input=params.transposed_input_2x3().repeat(bs, in_channels, 1, 1),
        oh=4,
        ow=5,
        sample_map=params.sample_map25(),
        cuda=True)

    # Manually computed correct result
    correct_output = 2 + params.in_channels * torch.tensor([[
        2, 0, 3, 0, 13
    ], [8, 0, 5, 6, 0], [8, 0, 8, 0, 4], [0, 0, 3, 0, 0]]).double().cuda()

    # Assert gradient check has passed
    assert gradcheck_res

    # Assert outputs match
    torch.testing.assert_close(output, correct_output)
//  [python]
def test_mean(self):
        ag = aggregators.MeanAggregator()
        torch.testing.assert_allclose(
            [[1.5]], ag(torch.Tensor([[1], [2]]), torch.LongTensor([0, 0]))
        )
        torch.testing.assert_allclose(
            [[3.5]], ag(torch.Tensor([[3], [4]]), torch.LongTensor([0, 0]))
        )
//  [python]
def test_mean(self):
        ag = aggregators.MeanAggregator()
        torch.testing.assert_close(
            [[1.5]], ag(torch.Tensor([[1], [2]]), torch.LongTensor([0, 0]))
        )
        torch.testing.assert_close(
            [[3.5]], ag(torch.Tensor([[3], [4]]), torch.LongTensor([0, 0]))
        )
//  [python]
def test_scatter(self):
        ag = aggregators.MeanAggregator()
        a = torch.Tensor([[1], [2], [3], [3], [4]])
        torch.testing.assert_allclose(
            ag.forward(a, torch.LongTensor([0, 0, 0, 1, 1])), [[2], [3.5]]
        )
//  [python]
def test_scatter(self):
        ag = aggregators.MeanAggregator()
        a = torch.Tensor([[1], [2], [3], [3], [4]])
        torch.testing.assert_close(
            ag.forward(a, torch.LongTensor([0, 0, 0, 1, 1])), [[2], [3.5]]
        )
//  [python]
def test_fold_sharded(device):
    torch.manual_seed(0)

    shape = (20, 230, 115, 8)
    N, H, W, C = shape
    stride_h = 2
    stride_w = 1

    torch_input = torch.randn(shape, dtype=torch.bfloat16)

    expected = fold_torch(torch_input, stride_h, stride_w)
    expected = expected.reshape(1, 1, -1, expected.shape[-1])

    shard_grid = tt.CoreRangeSet(
        {
            tt.CoreRange(
                tt.CoreCoord(0, 0),
                tt.CoreCoord(11, 7),
            ),
            tt.CoreRange(
                tt.CoreCoord(0, 8),
                tt.CoreCoord(3, 8),
            ),
        }
    )
    n_cores = 100

    shard_spec = tt.ShardSpec(shard_grid, [N * H * W // n_cores, C], tt.ShardOrientation.ROW_MAJOR, False)

    tt_input = torch2tt_tensor(
        torch_input,
        device,
        tt.Layout.ROW_MAJOR,
        tt_memory_config=tt.MemoryConfig(tt.TensorMemoryLayout.HEIGHT_SHARDED, tt.BufferType.L1, shard_spec),
    )
    tt_out = tt.fold(tt_input, stride_h, stride_w)
    actual = tt2torch_tensor(tt_out)

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_fold_sharded(device):
    torch.manual_seed(0)

    shape = (20, 230, 115, 8)
    N, H, W, C = shape
    stride_h = 2
    stride_w = 1

    torch_input = torch.randn(shape, dtype=torch.bfloat16)

    expected = fold_torch(torch_input, stride_h, stride_w)
    expected = expected.reshape(1, 1, -1, expected.shape[-1])

    shard_grid = tt.CoreRangeSet(
        {
            tt.CoreRange(
                tt.CoreCoord(0, 0),
                tt.CoreCoord(11, 7),
            ),
            tt.CoreRange(
                tt.CoreCoord(0, 8),
                tt.CoreCoord(3, 8),
            ),
        }
    )
    n_cores = 100

    shard_spec = tt.ShardSpec(shard_grid, [N * H * W // n_cores, C], tt.ShardOrientation.ROW_MAJOR, False)

    tt_input = torch2tt_tensor(
        torch_input,
        device,
        tt.Layout.ROW_MAJOR,
        tt_memory_config=tt.MemoryConfig(tt.TensorMemoryLayout.HEIGHT_SHARDED, tt.BufferType.L1, shard_spec),
    )
    tt_out = tt.fold(tt_input, stride_h, stride_w)
    actual = tt2torch_tensor(tt_out)

    torch.testing.assert_close(actual, expected)
//  [python]
def test_multidim_scatter(self):
        ag = aggregators.MeanAggregator()
        a = torch.Tensor([[[1], [2], [3], [3], [4]], [[1], [2], [3], [3], [4]]])
        indices = torch.LongTensor(
            [[[0], [0], [0], [1], [1]], [[1], [1], [1], [0], [0]]]
        )
        res = torch.tensor([[[2], [3.5]], [[3.5], [2]]])

        torch.testing.assert_allclose(ag.forward(a, indices, dim=1), res)
//  [python]
def test_multidim_scatter(self):
        ag = aggregators.MeanAggregator()
        a = torch.Tensor([[[1], [2], [3], [3], [4]], [[1], [2], [3], [3], [4]]])
        indices = torch.LongTensor(
            [[[0], [0], [0], [1], [1]], [[1], [1], [1], [0], [0]]]
        )
        res = torch.tensor([[[2], [3.5]], [[3.5], [2]]])

        torch.testing.assert_close(ag.forward(a, indices, dim=1), res)
//  [python]
def test_gather_uint32():
    from tetranerf.utils.extension import gather_uint32

    vals = torch.rand((5)).cuda()
    indices = torch.randint(0, 5, (12,), dtype=torch.int32).cuda()
    res = gather_uint32(vals, 0, indices).cuda()
    torch.testing.assert_allclose(res, vals[indices.long()])
    
    with pytest.raises(Exception):
        # Does not support dim > 1
        vals = torch.rand((5, 3)).cuda()
        indices = torch.randint(0, 3, (5, 8), dtype=torch.int32).cuda()
        gather_uint32(vals, 0, indices.long())
//  [python]
def test_gather_uint32():
    from tetranerf.utils.extension import gather_uint32

    vals = torch.rand((5)).cuda()
    indices = torch.randint(0, 5, (12,), dtype=torch.int32).cuda()
    res = gather_uint32(vals, 0, indices).cuda()
    torch.testing.assert_close(res, vals[indices.long()])
    
    with pytest.raises(Exception):
        # Does not support dim > 1
        vals = torch.rand((5, 3)).cuda()
        indices = torch.randint(0, 3, (5, 8), dtype=torch.int32).cuda()
        gather_uint32(vals, 0, indices.long())
//  [python]
def test_unscented_transform_julier_identity(dim: int):
    """Check unscented transform on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(dim=dim)
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_allclose(output_mean, expected_output_mean)
    torch.testing.assert_allclose(output_covariance, expected_output_covariance)
//  [python]
def test_unscented_transform_julier_identity(dim: int):
    """Check unscented transform on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(dim=dim)
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_close(output_mean, expected_output_mean)
    torch.testing.assert_close(output_covariance, expected_output_covariance)
//  [python]
def test_scatter_ema_uint32():
    from tetranerf.utils.extension import scatter_ema_uint32_

    torch.manual_seed(0)
    tensor = torch.rand((10)).cuda()
    indices = torch.tensor([4,3,5,8,2,1,0], dtype=torch.int32).cuda()
    vals = torch.rand((7,)).cuda()
    res = tensor.clone()
    decay = 0.5
    scatter_ema_uint32_(res, 0, indices, decay, vals)

    gt = torch.scatter(tensor, 0, indices.long(), tensor[indices.long()] * decay + (1-decay) * vals)
    torch.testing.assert_allclose(res, gt)
    
    with pytest.raises(Exception):
        # Does not support dim > 1
        tensor = torch.rand((5, 3)).cuda()
        indices = torch.randint(0, 3, (5, 8), dtype=torch.int32).cuda()
        vals = torch.rand((5, 8)).cuda()
        res = tensor.clone()
        scatter_ema_uint32_(res, 0, indices, decay, vals)
//  [python]
def test_scatter_ema_uint32():
    from tetranerf.utils.extension import scatter_ema_uint32_

    torch.manual_seed(0)
    tensor = torch.rand((10)).cuda()
    indices = torch.tensor([4,3,5,8,2,1,0], dtype=torch.int32).cuda()
    vals = torch.rand((7,)).cuda()
    res = tensor.clone()
    decay = 0.5
    scatter_ema_uint32_(res, 0, indices, decay, vals)

    gt = torch.scatter(tensor, 0, indices.long(), tensor[indices.long()] * decay + (1-decay) * vals)
    torch.testing.assert_close(res, gt)
    
    with pytest.raises(Exception):
        # Does not support dim > 1
        tensor = torch.rand((5, 3)).cuda()
        indices = torch.randint(0, 3, (5, 8), dtype=torch.int32).cuda()
        vals = torch.rand((5, 8)).cuda()
        res = tensor.clone()
        scatter_ema_uint32_(res, 0, indices, decay, vals)
//  [python]
def test_unscented_transform_julier_linear(dim: int):
    """Check unscented transform on linearly transformed multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Sample linear transformation matrix
    A = torch.randn(input_covariance.shape)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(dim=dim)
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    sigma_points = (A[:, None, :, :] @ sigma_points[:, :, :, None]).squeeze(-1)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = (A @ input_mean[:, :, None]).squeeze(-1)
    expected_output_covariance = A @ input_covariance @ A.transpose(-1, -2)
    torch.testing.assert_allclose(output_mean, expected_output_mean)
    torch.testing.assert_allclose(
        output_covariance, expected_output_covariance, atol=1e-4, rtol=1e-4
    )
//  [python]
def test_unscented_transform_julier_linear(dim: int):
    """Check unscented transform on linearly transformed multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Sample linear transformation matrix
    A = torch.randn(input_covariance.shape)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(dim=dim)
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    sigma_points = (A[:, None, :, :] @ sigma_points[:, :, :, None]).squeeze(-1)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = (A @ input_mean[:, :, None]).squeeze(-1)
    expected_output_covariance = A @ input_covariance @ A.transpose(-1, -2)
    torch.testing.assert_close(output_mean, expected_output_mean)
    torch.testing.assert_close(
        output_covariance, expected_output_covariance, atol=1e-4, rtol=1e-4
    )
//  [python]
def test_unscented_transform_merwe_identity(dim: int):
    """Check unscented transform on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_allclose(output_mean, expected_output_mean)
    torch.testing.assert_allclose(output_covariance, expected_output_covariance)
//  [python]
def test_unscented_transform_merwe_identity(dim: int):
    """Check unscented transform on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_close(output_mean, expected_output_mean)
    torch.testing.assert_close(output_covariance, expected_output_covariance)
//  [python]
def test_unscented_transform_merwe_linear(dim: int):
    """Check unscented transform on linearly transformed multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Sample linear transformation matrix
    A = torch.randn(input_covariance.shape)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    sigma_points = (A[:, None, :, :] @ sigma_points[:, :, :, None]).squeeze(-1)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = (A @ input_mean[:, :, None]).squeeze(-1)
    expected_output_covariance = A @ input_covariance @ A.transpose(-1, -2)
    torch.testing.assert_allclose(
        output_mean, expected_output_mean, atol=1e-4, rtol=1e-4
    )
    torch.testing.assert_allclose(
        output_covariance, expected_output_covariance, atol=1e-4, rtol=1e-4
    )
//  [python]
def test_unscented_transform_merwe_linear(dim: int):
    """Check unscented transform on linearly transformed multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Sample linear transformation matrix
    A = torch.randn(input_covariance.shape)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points(input_mean, input_covariance)
    sigma_points = (A[:, None, :, :] @ sigma_points[:, :, :, None]).squeeze(-1)
    output_mean, output_covariance = unscented_transform.compute_distribution(
        sigma_points
    )

    # Check outputs
    expected_output_mean = (A @ input_mean[:, :, None]).squeeze(-1)
    expected_output_covariance = A @ input_covariance @ A.transpose(-1, -2)
    torch.testing.assert_close(
        output_mean, expected_output_mean, atol=1e-4, rtol=1e-4
    )
    torch.testing.assert_close(
        output_covariance, expected_output_covariance, atol=1e-4, rtol=1e-4
    )
//  [python]
def test_yaw_to_quat(yaw_rad: Tensor) -> None:
    """Unit test converting yaw to a scalar first quaternion."""
    quats_wxyz = yaw_to_quat(yaw_rad)
    quats_wxyz_ = torch.as_tensor(
        R.from_euler("z", yaw_rad.numpy()).as_quat()
    )[..., [3, 0, 1, 2]]
    torch.testing.assert_allclose(quats_wxyz, quats_wxyz_)
//  [python]
def test_yaw_to_quat(yaw_rad: Tensor) -> None:
    """Unit test converting yaw to a scalar first quaternion."""
    quats_wxyz = yaw_to_quat(yaw_rad)
    quats_wxyz_ = torch.as_tensor(
        R.from_euler("z", yaw_rad.numpy()).as_quat()
    )[..., [3, 0, 1, 2]]
    torch.testing.assert_close(quats_wxyz, quats_wxyz_)
//  [python]
def test_xyz_to_quat(xyz_rad: Tensor) -> None:
    """Unit test converting Tait-Bryan angles to a scalar first quaternion.

    Args:
        xyz_rad: Roll, pitch, and yaw in radians.

    NOTE: The Tait-Bryan angles are equivalent to roll, pitch and yaw.
    """
    quats_wxyz = xyz_to_quat(xyz_rad)
    quats_wxyz_ = torch.as_tensor(
        R.from_euler("xyz", xyz_rad.numpy()).as_quat()
    )[..., [3, 0, 1, 2]]
    torch.testing.assert_allclose(quats_wxyz, quats_wxyz_)
//  [python]
def test_xyz_to_quat(xyz_rad: Tensor) -> None:
    """Unit test converting Tait-Bryan angles to a scalar first quaternion.

    Args:
        xyz_rad: Roll, pitch, and yaw in radians.

    NOTE: The Tait-Bryan angles are equivalent to roll, pitch and yaw.
    """
    quats_wxyz = xyz_to_quat(xyz_rad)
    quats_wxyz_ = torch.as_tensor(
        R.from_euler("xyz", xyz_rad.numpy()).as_quat()
    )[..., [3, 0, 1, 2]]
    torch.testing.assert_close(quats_wxyz, quats_wxyz_)
//  [python]
def test_unscented_transform_merwe_square_root(dim: int):
    """Check unscented transform square root formulation on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points_square_root(
        input_mean, torch.cholesky(input_covariance)
    )
    (
        output_mean,
        output_scale_tril,
    ) = unscented_transform.compute_distribution_square_root(sigma_points)

    # Check outputs
    output_covariance = output_scale_tril @ output_scale_tril.transpose(-1, -2)
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_allclose(output_mean, expected_output_mean)
    torch.testing.assert_allclose(output_covariance, expected_output_covariance)
//  [python]
def test_unscented_transform_merwe_square_root(dim: int):
    """Check unscented transform square root formulation on multivariate Gaussians."""
    input_mean, input_covariance = _gen_test_data(dim=dim)

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points_square_root(
        input_mean, torch.cholesky(input_covariance)
    )
    (
        output_mean,
        output_scale_tril,
    ) = unscented_transform.compute_distribution_square_root(sigma_points)

    # Check outputs
    output_covariance = output_scale_tril @ output_scale_tril.transpose(-1, -2)
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance
    torch.testing.assert_close(output_mean, expected_output_mean)
    torch.testing.assert_close(output_covariance, expected_output_covariance)
//  [python]
def test_simple(self, cuda=False):
        gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])
        proposed_instances = [
            torch.tensor([0, 2]),  # 100% instance 1
            torch.tensor([1, 4]),  # 2/3 of instance 2
            torch.tensor([3, 5]),  # 1/3 of instance 2 and 1/1 of instance 3
        ]
        if cuda:
            proposed_instances = [c.cuda() for c in proposed_instances]
            gt_instances = gt_instances.cuda()
        ious = instance_iou(proposed_instances, gt_instances)
        torch.testing.assert_allclose(
            ious.cpu(),
            torch.tensor([[1, 0, 0], [0, 2 / 3.0, 0], [0, 1.0 / 4.0, 1.0 / 2.0]]),
        )
//  [python]
def test_simple(self, cuda=False):
        gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])
        proposed_instances = [
            torch.tensor([0, 2]),  # 100% instance 1
            torch.tensor([1, 4]),  # 2/3 of instance 2
            torch.tensor([3, 5]),  # 1/3 of instance 2 and 1/1 of instance 3
        ]
        if cuda:
            proposed_instances = [c.cuda() for c in proposed_instances]
            gt_instances = gt_instances.cuda()
        ious = instance_iou(proposed_instances, gt_instances)
        torch.testing.assert_close(
            ious.cpu(),
            torch.tensor([[1, 0, 0], [0, 2 / 3.0, 0], [0, 1.0 / 4.0, 1.0 / 2.0]]),
        )
//  [python]
def test_unscented_transform_merwe_square_root_additive_nosie(dim: int):
    """Check unscented transform square root formulation on multivariate Gaussians, with
    additive noise.
    """
    input_mean, input_covariance = _gen_test_data(dim=dim)
    N, _dim = input_mean.shape
    assert _dim == dim

    additive_noise_scale_tril = torch.tril(torch.randn((N, dim, dim)))
    additive_noise_covariance = (
        additive_noise_scale_tril @ additive_noise_scale_tril.transpose(-1, -2)
    )

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points_square_root(
        input_mean, torch.cholesky(input_covariance)
    )
    (
        output_mean,
        output_scale_tril,
    ) = unscented_transform.compute_distribution_square_root(
        sigma_points,
        additive_noise_scale_tril=additive_noise_scale_tril,
    )

    # Check outputs
    output_covariance = output_scale_tril @ output_scale_tril.transpose(-1, -2)
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance + additive_noise_covariance
    torch.testing.assert_allclose(output_mean, expected_output_mean)
    torch.testing.assert_allclose(output_covariance, expected_output_covariance)
//  [python]
def test_unscented_transform_merwe_square_root_additive_nosie(dim: int):
    """Check unscented transform square root formulation on multivariate Gaussians, with
    additive noise.
    """
    input_mean, input_covariance = _gen_test_data(dim=dim)
    N, _dim = input_mean.shape
    assert _dim == dim

    additive_noise_scale_tril = torch.tril(torch.randn((N, dim, dim)))
    additive_noise_covariance = (
        additive_noise_scale_tril @ additive_noise_scale_tril.transpose(-1, -2)
    )

    # Perform unscented transform
    unscented_transform = torchfilter.utils.UnscentedTransform(
        dim=dim,
        sigma_point_strategy=torchfilter.utils.MerweSigmaPointStrategy(alpha=2e-1),
    )
    sigma_points = unscented_transform.select_sigma_points_square_root(
        input_mean, torch.cholesky(input_covariance)
    )
    (
        output_mean,
        output_scale_tril,
    ) = unscented_transform.compute_distribution_square_root(
        sigma_points,
        additive_noise_scale_tril=additive_noise_scale_tril,
    )

    # Check outputs
    output_covariance = output_scale_tril @ output_scale_tril.transpose(-1, -2)
    expected_output_mean = input_mean
    expected_output_covariance = input_covariance + additive_noise_covariance
    torch.testing.assert_close(output_mean, expected_output_mean)
    torch.testing.assert_close(output_covariance, expected_output_covariance)
//  [python]
def test_same(self):
        gt_instances = torch.randint(0, 10, (1000,))
        proposed_instances = []
        for i in range(20):
            instance_size = random.randint(5, 50)
            proposed_instances.append(torch.randint(0, 1000, (instance_size,)))

        ious = instance_iou(proposed_instances, gt_instances)
        proposed_instances_cuda = [i.cuda() for i in proposed_instances]
        ious_cuda = instance_iou(proposed_instances_cuda, gt_instances.cuda())

        torch.testing.assert_allclose(ious, ious_cuda.cpu())
//  [python]
def test_same(self):
        gt_instances = torch.randint(0, 10, (1000,))
        proposed_instances = []
        for i in range(20):
            instance_size = random.randint(5, 50)
            proposed_instances.append(torch.randint(0, 1000, (instance_size,)))

        ious = instance_iou(proposed_instances, gt_instances)
        proposed_instances_cuda = [i.cuda() for i in proposed_instances]
        ious_cuda = instance_iou(proposed_instances_cuda, gt_instances.cuda())

        torch.testing.assert_close(ious, ious_cuda.cpu())
//  [python]
def test_forward(self, batch_size, inner_size, feature_size, padding_value,
                     device):
        sizes = torch.randint(low=1,
                              high=inner_size,
                              size=(batch_size, ),
                              device=device)
        input_size = sizes.sum().item()
        x = torch.randn(input_size, device=device)
        y = self.split_and_pad(x, sizes, feature_size, padding_value)
        y_ref = split_and_pad_ref(x, sizes, feature_size, padding_value)
        torch.testing.assert_allclose(y, y_ref)
//  [python]
def test_forward(self, batch_size, inner_size, feature_size, padding_value,
                     device):
        sizes = torch.randint(low=1,
                              high=inner_size,
                              size=(batch_size, ),
                              device=device)
        input_size = sizes.sum().item()
        x = torch.randn(input_size, device=device)
        y = self.split_and_pad(x, sizes, feature_size, padding_value)
        y_ref = split_and_pad_ref(x, sizes, feature_size, padding_value)
        torch.testing.assert_close(y, y_ref)
//  [python]
def test_batch(self, cuda=False):
        gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])
        batch = torch.tensor([0, 0, 1, 1, 1, 1, 1])
        proposed_instances = [
            torch.tensor([0, 1]),  # 50% instance 1, 50% instance 2 of sample 1
            torch.tensor([3, 4]),  # 100% instance 2 of sample 2
            torch.tensor([5]),  # 100% of instance 3 of sample 2
        ]
        if cuda:
            proposed_instances = [c.cuda() for c in proposed_instances]
            gt_instances = gt_instances.cuda()
            batch = batch.cuda()
        ious = instance_iou(proposed_instances, gt_instances, batch=batch)
        torch.testing.assert_allclose(
            ious.cpu(),
            torch.tensor(
                [
                    [0.5, 0.5, 0, 0, 0],
                    [0, 0, 0, 1, 0],
                    [0, 0, 0, 0, 1],
                ]
            ),
        )
//  [python]
def test_batch(self, cuda=False):
        gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])
        batch = torch.tensor([0, 0, 1, 1, 1, 1, 1])
        proposed_instances = [
            torch.tensor([0, 1]),  # 50% instance 1, 50% instance 2 of sample 1
            torch.tensor([3, 4]),  # 100% instance 2 of sample 2
            torch.tensor([5]),  # 100% of instance 3 of sample 2
        ]
        if cuda:
            proposed_instances = [c.cuda() for c in proposed_instances]
            gt_instances = gt_instances.cuda()
            batch = batch.cuda()
        ious = instance_iou(proposed_instances, gt_instances, batch=batch)
        torch.testing.assert_close(
            ious.cpu(),
            torch.tensor(
                [
                    [0.5, 0.5, 0, 0, 0],
                    [0, 0, 0, 1, 0],
                    [0, 0, 0, 0, 1],
                ]
            ),
        )
//  [python]
def test_compute(self):
        npoints = 100
        pos = torch.randn((npoints, 3))
        x = torch.randn((npoints, 4))
        support = Data(x=x, pos=pos)
        query = Data(x=torch.randn((npoints // 2, 4)), pos=torch.randn((npoints // 2, 3)))

        interpolate = KNNInterpolate(3)
        precomputed = interpolate.precompute(query, support)

        gt = interpolate(query, support)
        pre = interpolate(query, support, precomputed=precomputed)

        torch.testing.assert_allclose(gt, pre)
//  [python]
def test_compute(self):
        npoints = 100
        pos = torch.randn((npoints, 3))
        x = torch.randn((npoints, 4))
        support = Data(x=x, pos=pos)
        query = Data(x=torch.randn((npoints // 2, 4)), pos=torch.randn((npoints // 2, 3)))

        interpolate = KNNInterpolate(3)
        precomputed = interpolate.precompute(query, support)

        gt = interpolate(query, support)
        pre = interpolate(query, support, precomputed=precomputed)

        torch.testing.assert_close(gt, pre)
//  [python]
def test_backward(self, batch_size, inner_size, device):
        sizes = torch.randint(low=1,
                              high=inner_size,
                              size=(batch_size, ),
                              device=device)
        input_size = sizes.sum().item()
        x = torch.randn(input_size, device=device)
        x_ref = x.detach().clone()
        x.requires_grad_(True)
        x_ref.requires_grad_(True)
        y = self.split_and_pad(x, sizes)
        y_ref = split_and_pad_ref(x_ref, sizes)

        dy = torch.randn_like(y)
        y.backward(dy)
        y_ref.backward(dy)
        torch.testing.assert_allclose(x.grad, x_ref.grad)
//  [python]
def test_backward(self, batch_size, inner_size, device):
        sizes = torch.randint(low=1,
                              high=inner_size,
                              size=(batch_size, ),
                              device=device)
        input_size = sizes.sum().item()
        x = torch.randn(input_size, device=device)
        x_ref = x.detach().clone()
        x.requires_grad_(True)
        x_ref.requires_grad_(True)
        y = self.split_and_pad(x, sizes)
        y_ref = split_and_pad_ref(x_ref, sizes)

        dy = torch.randn_like(y)
        y.backward(dy)
        y_ref.backward(dy)
        torch.testing.assert_close(x.grad, x_ref.grad)