jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_collate_custom_temporal_array_floats(pad_value):
    VOCAB_SIZE = 500

    cuts = CutSet.from_json("test/fixtures/ljspeech/cuts.json")
    max_num_frames = max(cut.num_frames for cut in cuts)

    with TemporaryDirectory() as d, NumpyFilesWriter(d) as writer:
        expected_posteriors = []
        for cut in cuts:
            expected_posteriors.append(
                np.random.randn(cut.num_frames, VOCAB_SIZE).astype(np.float32)
            )
            cut.posterior = writer.store_array(
                cut.id,
                expected_posteriors[-1],
                frame_shift=cut.frame_shift,
                temporal_dim=0,
            )

        posteriors, posterior_lens = collate_custom_field(
            cuts, "posterior", pad_value=pad_value
        )

        assert isinstance(posterior_lens, torch.Tensor)
        assert posterior_lens.dtype == torch.int32
        assert posterior_lens.shape == (len(cuts),)
        assert posterior_lens.tolist() == [c.num_frames for c in cuts]

        assert isinstance(posteriors, torch.Tensor)
        assert posteriors.dtype == torch.float32
        assert posteriors.shape == (len(cuts), max_num_frames, VOCAB_SIZE)
        for idx, post in enumerate(expected_posteriors):
            exp_len = post.shape[0]
            torch.testing.assert_allclose(posteriors[idx, :exp_len], post)
            expected_pad_value = 0 if pad_value is None else pad_value
            torch.testing.assert_allclose(
                posteriors[idx, exp_len:],
                expected_pad_value * torch.ones_like(posteriors[idx, exp_len:]),
            )
//  [python]
def test_collate_custom_temporal_array_floats(pad_value):
    VOCAB_SIZE = 500

    cuts = CutSet.from_json("test/fixtures/ljspeech/cuts.json")
    max_num_frames = max(cut.num_frames for cut in cuts)

    with TemporaryDirectory() as d, NumpyFilesWriter(d) as writer:
        expected_posteriors = []
        for cut in cuts:
            expected_posteriors.append(
                np.random.randn(cut.num_frames, VOCAB_SIZE).astype(np.float32)
            )
            cut.posterior = writer.store_array(
                cut.id,
                expected_posteriors[-1],
                frame_shift=cut.frame_shift,
                temporal_dim=0,
            )

        posteriors, posterior_lens = collate_custom_field(
            cuts, "posterior", pad_value=pad_value
        )

        assert isinstance(posterior_lens, torch.Tensor)
        assert posterior_lens.dtype == torch.int32
        assert posterior_lens.shape == (len(cuts),)
        assert posterior_lens.tolist() == [c.num_frames for c in cuts]

        assert isinstance(posteriors, torch.Tensor)
        assert posteriors.dtype == torch.float32
        assert posteriors.shape == (len(cuts), max_num_frames, VOCAB_SIZE)
        for idx, post in enumerate(expected_posteriors):
            exp_len = post.shape[0]
            torch.testing.assert_close(posteriors[idx, :exp_len], post)
            expected_pad_value = 0 if pad_value is None else pad_value
            torch.testing.assert_close(
                posteriors[idx, exp_len:],
                expected_pad_value * torch.ones_like(posteriors[idx, exp_len:]),
            )
//  [python]
def test_set_gt_score(self):
        data_sample = ClsDataSample()
        data_sample.set_gt_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('score', data_sample.gt_label)
        torch.testing.assert_allclose(data_sample.gt_label.score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test set again
        data_sample.set_gt_score(torch.tensor([0.2, 0.1, 0.5, 0.1, 0.1]))
        torch.testing.assert_allclose(data_sample.gt_label.score,
                                      [0.2, 0.1, 0.5, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_gt_score(torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_set_gt_score(self):
        data_sample = ClsDataSample()
        data_sample.set_gt_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('score', data_sample.gt_label)
        torch.testing.assert_close(data_sample.gt_label.score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test set again
        data_sample.set_gt_score(torch.tensor([0.2, 0.1, 0.5, 0.1, 0.1]))
        torch.testing.assert_close(data_sample.gt_label.score,
                                      [0.2, 0.1, 0.5, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_gt_score(torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_set_pred_score(self):
        data_sample = ClsDataSample()
        data_sample.set_pred_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('score', data_sample.pred_label)
        torch.testing.assert_allclose(data_sample.pred_label.score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test set again
        data_sample.set_pred_score(torch.tensor([0.2, 0.1, 0.5, 0.1, 0.1]))
        torch.testing.assert_allclose(data_sample.pred_label.score,
                                      [0.2, 0.1, 0.5, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_pred_score(
                torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_set_pred_score(self):
        data_sample = ClsDataSample()
        data_sample.set_pred_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('score', data_sample.pred_label)
        torch.testing.assert_close(data_sample.pred_label.score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test set again
        data_sample.set_pred_score(torch.tensor([0.2, 0.1, 0.5, 0.1, 0.1]))
        torch.testing.assert_close(data_sample.pred_label.score,
                                      [0.2, 0.1, 0.5, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_pred_score(
                torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_empty_layer():
    x = torch.randn(10, 8)
    model = Head()
    y = model(x)
    torch.testing.assert_allclose(x, y)
//  [python]
def test_empty_layer():
    x = torch.randn(10, 8)
    model = Head()
    y = model(x)
    torch.testing.assert_close(x, y)
//  [python]
def test_superpos(self):
        stack = MultiPushStack.empty(1, 3)
        stack.update(PUSH2, NEW_VEC)
        policy = torch.tensor([[1/2, 1/2, 0, 0, 0, 0]])
        stack.update(policy, NEW_VEC)

        # Should have .5 * [1] + .5 * [1, 1] = [1, .5].
        expected = [[
            [1., 1., 0.],
            [.5, .5, 0.],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
        ]]
        torch.testing.assert_allclose(stack.tapes.tolist(), expected)
//  [python]
def test_superpos(self):
        stack = MultiPushStack.empty(1, 3)
        stack.update(PUSH2, NEW_VEC)
        policy = torch.tensor([[1/2, 1/2, 0, 0, 0, 0]])
        stack.update(policy, NEW_VEC)

        # Should have .5 * [1] + .5 * [1, 1] = [1, .5].
        expected = [[
            [1., 1., 0.],
            [.5, .5, 0.],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
        ]]
        torch.testing.assert_close(stack.tapes.tolist(), expected)
//  [python]
def test_last_pooling_with_mask():
    lp = pooling.LastPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [5, 6, 7, 8],
            [9, 10, 11, 12]
        ]
    )
    torch.testing.assert_allclose(out, expected)
//  [python]
def test_last_pooling_with_mask():
    lp = pooling.LastPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [5, 6, 7, 8],
            [9, 10, 11, 12]
        ]
    )
    torch.testing.assert_close(out, expected)
//  [python]
def test_first_pooling_with_mask():
    lp = pooling.FirstPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [1, 2, 3, 4],
            [1, 2, 3, 4]
        ]
    )
    torch.testing.assert_allclose(out, expected)
//  [python]
def test_first_pooling_with_mask():
    lp = pooling.FirstPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [1, 2, 3, 4],
            [1, 2, 3, 4]
        ]
    )
    torch.testing.assert_close(out, expected)
//  [python]
def test_pca():
    pca = torchdrift.reducers.PCAReducer(n_components=2)
    assert "n_components" in str(pca)
    a = torch.randn(100, 50, dtype=torch.double)
    red = pca.fit(a)
    pca_ref = sklearn.decomposition.PCA(n_components=2)
    red_ref = torch.from_numpy(pca_ref.fit_transform(a))
    # need to find a way to deal with signs
    torch.testing.assert_allclose(red.abs(), red_ref.abs())
    b = torch.randn(25, 50, dtype=torch.double)
    red2 = pca(b)
    red2_ref = torch.from_numpy(pca_ref.transform(b))
//  [python]
def test_pca():
    pca = torchdrift.reducers.PCAReducer(n_components=2)
    assert "n_components" in str(pca)
    a = torch.randn(100, 50, dtype=torch.double)
    red = pca.fit(a)
    pca_ref = sklearn.decomposition.PCA(n_components=2)
    red_ref = torch.from_numpy(pca_ref.fit_transform(a))
    # need to find a way to deal with signs
    torch.testing.assert_close(red.abs(), red_ref.abs())
    b = torch.randn(25, 50, dtype=torch.double)
    red2 = pca(b)
    red2_ref = torch.from_numpy(pca_ref.transform(b))
//  [python]
def test_avg_pooling_with_mask():
    lp = pooling.AvgPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [3, 4, 5, 6],
            [5, 6, 7, 8]
        ]
    )
    torch.testing.assert_allclose(out, expected)
//  [python]
def test_avg_pooling_with_mask():
    lp = pooling.AvgPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [3, 4, 5, 6],
            [5, 6, 7, 8]
        ]
    )
    torch.testing.assert_close(out, expected)
//  [python]
def test_sum_pooling_with_mask():
    lp = pooling.SumPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [6, 8, 10, 12],
            [15, 18, 21, 24]
        ]
    )
    torch.testing.assert_allclose(out, expected)
//  [python]
def test_sum_pooling_with_mask():
    lp = pooling.SumPooling()
    out = lp(TENSOR, padding_mask=MASK)

    assert out.size() == torch.Size([2, 4])

    expected = Tensor(
        [
            [6, 8, 10, 12],
            [15, 18, 21, 24]
        ]
    )
    torch.testing.assert_close(out, expected)
//  [python]
def test_consistency(device, levy_area_approximation):
    if device == gpu and not torch.cuda.is_available():
        pytest.skip(msg="CUDA not available.")

    t0, t1 = 0.0, 1.0
    for _ in range(REPS):
        bm = torchsde.BrownianInterval(t0=t0, t1=t1, size=(LARGE_BATCH_SIZE,), device=device,
                                       levy_area_approximation=levy_area_approximation)

        for _ in range(MEDIUM_REPS):
            ta, t_, tb = sorted(npr.uniform(low=t0, high=t1, size=(3,)))

            if levy_area_approximation == 'none':
                W = bm(ta, tb)
                W1 = bm(ta, t_)
                W2 = bm(t_, tb)
            else:
                W, U = bm(ta, tb, return_U=True)
                W1, U1 = bm(ta, t_, return_U=True)
                W2, U2 = bm(t_, tb, return_U=True)

            torch.testing.assert_allclose(W1 + W2, W, rtol=1e-6, atol=1e-6)
            if levy_area_approximation != 'none':
                torch.testing.assert_allclose(U1 + U2 + (tb - t_) * W1, U, rtol=1e-6, atol=1e-6)
//  [python]
def test_consistency(device, levy_area_approximation):
    if device == gpu and not torch.cuda.is_available():
        pytest.skip(msg="CUDA not available.")

    t0, t1 = 0.0, 1.0
    for _ in range(REPS):
        bm = torchsde.BrownianInterval(t0=t0, t1=t1, size=(LARGE_BATCH_SIZE,), device=device,
                                       levy_area_approximation=levy_area_approximation)

        for _ in range(MEDIUM_REPS):
            ta, t_, tb = sorted(npr.uniform(low=t0, high=t1, size=(3,)))

            if levy_area_approximation == 'none':
                W = bm(ta, tb)
                W1 = bm(ta, t_)
                W2 = bm(t_, tb)
            else:
                W, U = bm(ta, tb, return_U=True)
                W1, U1 = bm(ta, t_, return_U=True)
                W2, U2 = bm(t_, tb, return_U=True)

            torch.testing.assert_close(W1 + W2, W, rtol=1e-6, atol=1e-6)
            if levy_area_approximation != 'none':
                torch.testing.assert_close(U1 + U2 + (tb - t_) * W1, U, rtol=1e-6, atol=1e-6)
//  [python]
def test_simple_gpu(self):
        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).to(torch.float).cuda()
        b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float).cuda()
        idx, dist = ball_query(1.01, 2, a, b)
        torch.testing.assert_allclose(idx.cpu(), torch.tensor([[[0, 1]], [[2, 2]]]))
        torch.testing.assert_allclose(dist.cpu(), torch.tensor([[[0, 1]], [[1, -1]]]).float())
//  [python]
def test_simple_gpu(self):
        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).to(torch.float).cuda()
        b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float).cuda()
        idx, dist = ball_query(1.01, 2, a, b)
        torch.testing.assert_close(idx.cpu(), torch.tensor([[[0, 1]], [[2, 2]]]))
        torch.testing.assert_close(dist.cpu(), torch.tensor([[[0, 1]], [[1, -1]]]).float())
//  [python]
def test_simple_cpu(self):
        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).to(torch.float)
        b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float)
        idx, dist = ball_query(1.01, 2, a, b, sort=True)
        torch.testing.assert_allclose(idx, torch.tensor([[[0, 1]], [[2, 2]]]))
        torch.testing.assert_allclose(dist, torch.tensor([[[0, 1]], [[1, -1]]]).float())

        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0]]]).to(torch.float)
        idx, dist = ball_query(1.01, 3, a, a, sort=True)
        torch.testing.assert_allclose(idx, torch.tensor([[[0, 1, 0], [1, 0, 2], [2, 1, 2]]]))
//  [python]
def test_simple_cpu(self):
        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).to(torch.float)
        b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float)
        idx, dist = ball_query(1.01, 2, a, b, sort=True)
        torch.testing.assert_close(idx, torch.tensor([[[0, 1]], [[2, 2]]]))
        torch.testing.assert_close(dist, torch.tensor([[[0, 1]], [[1, -1]]]).float())

        a = torch.tensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0]]]).to(torch.float)
        idx, dist = ball_query(1.01, 3, a, a, sort=True)
        torch.testing.assert_close(idx, torch.tensor([[[0, 1, 0], [1, 0, 2], [2, 1, 2]]]))
//  [python]
def test_to_copy(self, dtype: torch.dtype):
        inpt_tensor = torch.rand(128, device='cpu')
        inpt_tensor_nf4 = to_nf4(inpt_tensor, 32, 2)
        nf4_to_dtype = inpt_tensor_nf4.to(dtype)
        torch.testing.assert_allclose(inpt_tensor, nf4_to_dtype, atol=0.13, rtol=0.13)

        if torch.cuda.is_available():
            inpt_tensor = torch.rand(128, device='cuda')
            inpt_tensor_nf4 = to_nf4(inpt_tensor, 32, 2)
            nf4_to_dtype = inpt_tensor_nf4.to(dtype)
            torch.testing.assert_allclose(inpt_tensor, nf4_to_dtype, atol=0.13, rtol=0.13)
//  [python]
def test_to_copy(self, dtype: torch.dtype):
        inpt_tensor = torch.rand(128, device='cpu')
        inpt_tensor_nf4 = to_nf4(inpt_tensor, 32, 2)
        nf4_to_dtype = inpt_tensor_nf4.to(dtype)
        torch.testing.assert_close(inpt_tensor, nf4_to_dtype, atol=0.13, rtol=0.13)

        if torch.cuda.is_available():
            inpt_tensor = torch.rand(128, device='cuda')
            inpt_tensor_nf4 = to_nf4(inpt_tensor, 32, 2)
            nf4_to_dtype = inpt_tensor_nf4.to(dtype)
            torch.testing.assert_close(inpt_tensor, nf4_to_dtype, atol=0.13, rtol=0.13)
//  [python]
def test_discrete(t_start, q_function):
    policy = SoftMax(q_function, t_start)
    for _ in range(100):
        state = torch.randint(4, ())
        logits = q_function(state)[..., 0]
        probs = torch.softmax(logits / t_start, dim=0)
        torch.testing.assert_allclose(
            tensor_to_distribution(policy(state)).probs, probs
        )
//  [python]
def test_discrete(t_start, q_function):
    policy = SoftMax(q_function, t_start)
    for _ in range(100):
        state = torch.randint(4, ())
        logits = q_function(state)[..., 0]
        probs = torch.softmax(logits / t_start, dim=0)
        torch.testing.assert_close(
            tensor_to_distribution(policy(state)).probs, probs
        )
//  [python]
def test_basic(self):

        shape = (512, 512, 3)
        x = torch.randn(512, 512, 3)
        y = torch.randn(512, 512, 3)
        z = torch.randn(512, 512, 3)

        @torch.jit.script
        def foo(a, b, c):
            return a * b + c

        jit_out, pml_out = self._run_both(foo, [x, y, z])
        torch.testing.assert_allclose(jit_out, pml_out, rtol=0.01, atol=0.01)
//  [python]
def test_basic(self):

        shape = (512, 512, 3)
        x = torch.randn(512, 512, 3)
        y = torch.randn(512, 512, 3)
        z = torch.randn(512, 512, 3)

        @torch.jit.script
        def foo(a, b, c):
            return a * b + c

        jit_out, pml_out = self._run_both(foo, [x, y, z])
        torch.testing.assert_close(jit_out, pml_out, rtol=0.01, atol=0.01)
//  [python]
def test_superpos(self):
        stack = MinimalistStack.empty(1, 3)
        stack.update(PUSH, VEC1)
        stack.update(PUSH, VEC1)
        policy = torch.ones(1, 2) / 2
        stack.update(policy, VEC1)
        expected = [[[1., 1., 0], [1/2, 1/2, 0], [1/2, 1/2, 0]]]
        torch.testing.assert_allclose(stack.tapes.tolist(), expected)
//  [python]
def test_superpos(self):
        stack = MinimalistStack.empty(1, 3)
        stack.update(PUSH, VEC1)
        stack.update(PUSH, VEC1)
        policy = torch.ones(1, 2) / 2
        stack.update(policy, VEC1)
        expected = [[[1., 1., 0], [1/2, 1/2, 0], [1/2, 1/2, 0]]]
        torch.testing.assert_close(stack.tapes.tolist(), expected)
//  [python]
def test_packer():
    reset = torch.Tensor([
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]]).bool().cuda()

    x = torch.Tensor([
        [0, 10, 20, 30],
        [1, 11, 40, 31],
        [2, 12, 41, 50]]).float().cuda()

    h = torch.Tensor([[1, 2, 3, 4]]).float().cuda()

    packer = Packer(reset)
    xp, (hp, cp) = packer.pack(x, h, h)
    xu, (hu, cu) = packer.unpack(xp, (hp, cp))

    torch.testing.assert_allclose(xp.data, torch.tensor([0, 10, 20, 40, 30, 50, 1, 11, 41, 31, 2, 12]).float().cuda())
    torch.testing.assert_allclose(xp.batch_sizes, torch.tensor([6, 4, 2]))
    torch.testing.assert_allclose(x, xu)

    torch.testing.assert_allclose(hp, torch.tensor([1, 0, 3, 0, 4, 0]).float().cuda())
    torch.testing.assert_allclose(hu, torch.tensor([[1, 0, 0, 0]]).float().cuda())
//  [python]
def test_packer():
    reset = torch.Tensor([
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]]).bool().cuda()

    x = torch.Tensor([
        [0, 10, 20, 30],
        [1, 11, 40, 31],
        [2, 12, 41, 50]]).float().cuda()

    h = torch.Tensor([[1, 2, 3, 4]]).float().cuda()

    packer = Packer(reset)
    xp, (hp, cp) = packer.pack(x, h, h)
    xu, (hu, cu) = packer.unpack(xp, (hp, cp))

    torch.testing.assert_close(xp.data, torch.tensor([0, 10, 20, 40, 30, 50, 1, 11, 41, 31, 2, 12]).float().cuda())
    torch.testing.assert_close(xp.batch_sizes, torch.tensor([6, 4, 2]))
    torch.testing.assert_close(x, xu)

    torch.testing.assert_close(hp, torch.tensor([1, 0, 3, 0, 4, 0]).float().cuda())
    torch.testing.assert_close(hu, torch.tensor([[1, 0, 0, 0]]).float().cuda())
//  [python]
def test_measurement_jacobian():
    """Checks that our autograd-computed measurement jacobian is correct."""
    N = 10
    measurement_model = LinearKalmanFilterMeasurementModel()
    C_autograd = measurement_model.jacobian(states=torch.zeros((N, state_dim)))

    for i in range(N):
        torch.testing.assert_allclose(C_autograd[i], C)
//  [python]
def test_measurement_jacobian():
    """Checks that our autograd-computed measurement jacobian is correct."""
    N = 10
    measurement_model = LinearKalmanFilterMeasurementModel()
    C_autograd = measurement_model.jacobian(states=torch.zeros((N, state_dim)))

    for i in range(N):
        torch.testing.assert_close(C_autograd[i], C)
//  [python]
def test_dynamics_jacobian():
    """Checks that our autograd-computed dynamics jacobian is correct."""
    N = 10
    dynamics_model = LinearDynamicsModel()
    A_autograd = dynamics_model.jacobian(
        initial_states=torch.zeros((N, state_dim)),
        controls=torch.zeros((N, control_dim)),
    )

    for i in range(N):
        torch.testing.assert_allclose(A_autograd[i], A)
//  [python]
def test_dynamics_jacobian():
    """Checks that our autograd-computed dynamics jacobian is correct."""
    N = 10
    dynamics_model = LinearDynamicsModel()
    A_autograd = dynamics_model.jacobian(
        initial_states=torch.zeros((N, state_dim)),
        controls=torch.zeros((N, control_dim)),
    )

    for i in range(N):
        torch.testing.assert_close(A_autograd[i], A)
//  [python]
def assert_eq(actual, expected, msg="", rtol=None, atol=None):
    """Asserts two things are equal with nice support for lists and tensors

    It also gives prettier error messages than assert a == b
    """
    if not msg:
        msg = f"Values are not equal: \n\ta={actual} \n\tb={expected}"

    if isinstance(actual, torch.Size):
        actual = list(actual)
    if isinstance(expected, torch.Size):
        expected = list(expected)

    if isinstance(actual, tuple):
        actual = list(actual)
    if isinstance(expected, tuple):
        expected = list(expected)

    if isinstance(actual, torch.Tensor):
        if rtol is None and atol is None:
            rtol, atol = _get_default_rtol_and_atol(actual=actual, expected=expected)
        torch.testing.assert_allclose(actual, expected, msg=msg, rtol=rtol, atol=atol)
        return
    if isinstance(actual, np.ndarray):
        np.testing.assert_allclose(actual, expected, rtol=rtol or 0, atol=atol or 0)
        return
    if isinstance(actual, torch.Size) or isinstance(expected, torch.Size):
        assert actual == expected, msg
        return
    if isinstance(actual, dict):
        assert isinstance(expected, dict)
        assert actual.keys() == expected.keys(), msg
        for key in actual.keys():
            assert_eq(actual[key], expected[key], msg=msg, rtol=rtol, atol=atol)
        return
    if isinstance(actual, (tuple, list, set)):
        assert isinstance(expected, type(actual))
        assert len(actual) == len(expected), msg
        for ai, bi in zip(actual, expected):
            assert_eq(ai, bi, msg=msg, rtol=rtol, atol=atol)
        return

    if rtol is None and atol is None:
        assert actual == expected, f"{actual} != {expected}"
    else:
        atol = 0 if atol is None else atol
        rtol = 0 if rtol is None else rtol
        assert (
            abs(actual - expected) <= atol + expected * rtol
        ), f"{actual} != {expected}"
//  [python]
def assert_eq(actual, expected, msg="", rtol=None, atol=None):
    """Asserts two things are equal with nice support for lists and tensors

    It also gives prettier error messages than assert a == b
    """
    if not msg:
        msg = f"Values are not equal: \n\ta={actual} \n\tb={expected}"

    if isinstance(actual, torch.Size):
        actual = list(actual)
    if isinstance(expected, torch.Size):
        expected = list(expected)

    if isinstance(actual, tuple):
        actual = list(actual)
    if isinstance(expected, tuple):
        expected = list(expected)

    if isinstance(actual, torch.Tensor):
        if rtol is None and atol is None:
            rtol, atol = _get_default_rtol_and_atol(actual=actual, expected=expected)
        torch.testing.assert_close(actual, expected, msg=msg, rtol=rtol, atol=atol)
        return
    if isinstance(actual, np.ndarray):
        np.testing.assert_allclose(actual, expected, rtol=rtol or 0, atol=atol or 0)
        return
    if isinstance(actual, torch.Size) or isinstance(expected, torch.Size):
        assert actual == expected, msg
        return
    if isinstance(actual, dict):
        assert isinstance(expected, dict)
        assert actual.keys() == expected.keys(), msg
        for key in actual.keys():
            assert_eq(actual[key], expected[key], msg=msg, rtol=rtol, atol=atol)
        return
    if isinstance(actual, (tuple, list, set)):
        assert isinstance(expected, type(actual))
        assert len(actual) == len(expected), msg
        for ai, bi in zip(actual, expected):
            assert_eq(ai, bi, msg=msg, rtol=rtol, atol=atol)
        return

    if rtol is None and atol is None:
        assert actual == expected, f"{actual} != {expected}"
    else:
        atol = 0 if atol is None else atol
        rtol = 0 if rtol is None else rtol
        assert (
            abs(actual - expected) <= atol + expected * rtol
        ), f"{actual} != {expected}"
//  [python]
def test_reset_parameters(self):
        """Test the reset of the parameters."""
        node_channels = 6
        edge_channels = 10
        face_channels = 7

        scconv = SCConvLayer(node_channels, edge_channels, face_channels)
        scconv.reset_parameters()

        for module in scconv.modules():
            if isinstance(module, torch.nn.Conv2d):
                torch.testing.assert_allclose(
                    module.weight, torch.zeros_like(module.weight)
                )
                torch.testing.assert_allclose(
                    module.bias, torch.zeros_like(module.bias)
                )
//  [python]
def test_reset_parameters(self):
        """Test the reset of the parameters."""
        node_channels = 6
        edge_channels = 10
        face_channels = 7

        scconv = SCConvLayer(node_channels, edge_channels, face_channels)
        scconv.reset_parameters()

        for module in scconv.modules():
            if isinstance(module, torch.nn.Conv2d):
                torch.testing.assert_close(
                    module.weight, torch.zeros_like(module.weight)
                )
                torch.testing.assert_close(
                    module.bias, torch.zeros_like(module.bias)
                )
//  [python]
def test_random_cpu(self, cuda=False):
        a = torch.randn(100, 3).to(torch.float)
        b = torch.randn(50, 3).to(torch.float)
        batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
        batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
        R = 1

        idx, dist = ball_query(
            R,
            15,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        idx1, dist = ball_query(
            R,
            15,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        torch.testing.assert_allclose(idx1, idx)
        with self.assertRaises(AssertionError):
            idx, dist = ball_query(
                R,
                15,
                a,
                b,
                mode="PARTIAL_DENSE",
                batch_x=batch_a,
                batch_y=batch_b,
                sort=False,
            )
            idx1, dist = ball_query(
                R,
                15,
                a,
                b,
                mode="PARTIAL_DENSE",
                batch_x=batch_a,
                batch_y=batch_b,
                sort=False,
            )
            torch.testing.assert_allclose(idx1, idx)

        self.assertEqual(idx.shape[0], b.shape[0])
        self.assertEqual(dist.shape[0], b.shape[0])
        self.assertLessEqual(idx.max().item(), len(batch_a))

        # Comparison to see if we have the same result
        tree = KDTree(a.detach().numpy())
        idx3_sk = tree.query_radius(b.detach().numpy(), r=R)
        i = np.random.randint(len(batch_b))
        for p in idx[i].detach().numpy():
            if p >= 0 and p < len(batch_a):
                assert p in idx3_sk[i]
//  [python]
def test_random_cpu(self, cuda=False):
        a = torch.randn(100, 3).to(torch.float)
        b = torch.randn(50, 3).to(torch.float)
        batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
        batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
        R = 1

        idx, dist = ball_query(
            R,
            15,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        idx1, dist = ball_query(
            R,
            15,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        torch.testing.assert_close(idx1, idx)
        with self.assertRaises(AssertionError):
            idx, dist = ball_query(
                R,
                15,
                a,
                b,
                mode="PARTIAL_DENSE",
                batch_x=batch_a,
                batch_y=batch_b,
                sort=False,
            )
            idx1, dist = ball_query(
                R,
                15,
                a,
                b,
                mode="PARTIAL_DENSE",
                batch_x=batch_a,
                batch_y=batch_b,
                sort=False,
            )
            torch.testing.assert_close(idx1, idx)

        self.assertEqual(idx.shape[0], b.shape[0])
        self.assertEqual(dist.shape[0], b.shape[0])
        self.assertLessEqual(idx.max().item(), len(batch_a))

        # Comparison to see if we have the same result
        tree = KDTree(a.detach().numpy())
        idx3_sk = tree.query_radius(b.detach().numpy(), r=R)
        i = np.random.randint(len(batch_b))
        for p in idx[i].detach().numpy():
            if p >= 0 and p < len(batch_a):
                assert p in idx3_sk[i]
//  [python]
def test_conv_layer_wrapper(conv_input, conv_weights, mode, clip_acts, per_channel_wts, conv_stats, expected_output):
    layer = torch.nn.Conv2d(conv_input.shape[1], expected_output.shape[1], conv_weights.shape[-1],
                            padding=1, bias=False)
    layer.weight.data = conv_weights

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts, activation_stats=conv_stats)

    input_stats = None if conv_stats is None else conv_stats['inputs'][0]
    conv_input = attach_quant_metadata(conv_input, 8, mode, stats=input_stats, clip_mode=clip_acts,
                                       per_channel=False, num_stds=None, scale_approx_mult_bits=None)

    with pytest.raises(RuntimeError):
        model(conv_input)

    model.eval()

    output = model(conv_input)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_conv_layer_wrapper(conv_input, conv_weights, mode, clip_acts, per_channel_wts, conv_stats, expected_output):
    layer = torch.nn.Conv2d(conv_input.shape[1], expected_output.shape[1], conv_weights.shape[-1],
                            padding=1, bias=False)
    layer.weight.data = conv_weights

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts, activation_stats=conv_stats)

    input_stats = None if conv_stats is None else conv_stats['inputs'][0]
    conv_input = attach_quant_metadata(conv_input, 8, mode, stats=input_stats, clip_mode=clip_acts,
                                       per_channel=False, num_stds=None, scale_approx_mult_bits=None)

    with pytest.raises(RuntimeError):
        model(conv_input)

    model.eval()

    output = model(conv_input)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_linear_layer_wrapper(linear_input, linear_weights, linear_bias,
                              mode, clip_acts, per_channel_wts, expected_output):
    layer = torch.nn.Linear(linear_input.shape[1], expected_output.shape[1], bias=True)
    layer.weight.data = linear_weights
    layer.bias.data = linear_bias

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts)

    linear_input = attach_quant_metadata(linear_input, 8, mode, stats=None, clip_mode=clip_acts,
                                         per_channel=False, num_stds=None, scale_approx_mult_bits=None)

    # with pytest.raises(RuntimeError):
    #     model(linear_input)

    model.eval()

    output = model(linear_input)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_linear_layer_wrapper(linear_input, linear_weights, linear_bias,
                              mode, clip_acts, per_channel_wts, expected_output):
    layer = torch.nn.Linear(linear_input.shape[1], expected_output.shape[1], bias=True)
    layer.weight.data = linear_weights
    layer.bias.data = linear_bias

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts)

    linear_input = attach_quant_metadata(linear_input, 8, mode, stats=None, clip_mode=clip_acts,
                                         per_channel=False, num_stds=None, scale_approx_mult_bits=None)

    # with pytest.raises(RuntimeError):
    #     model(linear_input)

    model.eval()

    output = model(linear_input)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_concat_layer_wrapper(inputs, concat_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantConcatWrapper(torch.nn.Module(), 8, mode, clip_acts, concat_stats)

    layer = distiller.modules.Concat(dim=1)

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=concat_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, concat_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_concat_layer_wrapper(inputs, concat_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantConcatWrapper(torch.nn.Module(), 8, mode, clip_acts, concat_stats)

    layer = distiller.modules.Concat(dim=1)

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=concat_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, concat_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_eltwise_mult_layer_wrapper(inputs, eltwise_mult_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseMultWrapper(torch.nn.Module(), 8, mode, clip_acts, eltwise_mult_stats)

    layer = distiller.modules.EltwiseMult()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=eltwise_mult_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, eltwise_mult_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_eltwise_mult_layer_wrapper(inputs, eltwise_mult_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseMultWrapper(torch.nn.Module(), 8, mode, clip_acts, eltwise_mult_stats)

    layer = distiller.modules.EltwiseMult()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=eltwise_mult_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, eltwise_mult_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_eltwise_add_layer_wrapper(inputs, eltwise_add_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseAddWrapper(torch.nn.Module(), 8, mode, clip_acts, test_eltwise_add_layer_wrapper)

    layer = distiller.modules.EltwiseAdd()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=eltwise_add_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, eltwise_add_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_eltwise_add_layer_wrapper(inputs, eltwise_add_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseAddWrapper(torch.nn.Module(), 8, mode, clip_acts, test_eltwise_add_layer_wrapper)

    layer = distiller.modules.EltwiseAdd()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    for idx in range(len(inputs)):
        inputs[idx] = attach_quant_metadata(inputs[idx], 8, mode, stats=eltwise_add_stats['inputs'][idx],
                                            clip_mode=clip_acts, per_channel=False, num_stds=None,
                                            scale_approx_mult_bits=None)

    model = RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, eltwise_add_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_set_gt_score(self):
        data_sample = DataSample()
        data_sample.set_gt_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('gt_score', data_sample)
        torch.testing.assert_allclose(data_sample.gt_score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_gt_score(torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_set_gt_score(self):
        data_sample = DataSample()
        data_sample.set_gt_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('gt_score', data_sample)
        torch.testing.assert_close(data_sample.gt_score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_gt_score(torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_get_tensor_mean_n_stds_min_max():
    with pytest.raises(ValueError, message='Expecting ValueError with n_stds = 0'):
        qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=0)

    mean = torch.tensor(-16.)
    std = torch.tensor(62.87447738647461)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor)
    torch.testing.assert_allclose(t_min, mean - std)
    torch.testing.assert_allclose(t_max, mean + std)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=2)
    torch.testing.assert_allclose(t_min, torch.tensor(-95.))
    torch.testing.assert_allclose(t_max, torch.tensor(87.))
//  [python]
def test_get_tensor_mean_n_stds_min_max():
    with pytest.raises(ValueError, message='Expecting ValueError with n_stds = 0'):
        qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=0)

    mean = torch.tensor(-16.)
    std = torch.tensor(62.87447738647461)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor)
    torch.testing.assert_close(t_min, mean - std)
    torch.testing.assert_close(t_max, mean + std)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=2)
    torch.testing.assert_close(t_min, torch.tensor(-95.))
    torch.testing.assert_close(t_max, torch.tensor(87.))
//  [python]
def test_mag_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_mag_mask = torch.ones((1, fb_dim // 2, 1))
        masked = transforms.apply_mag_mask(tf_rep, id_mag_mask, dim=1)
        assert_allclose(masked, tf_rep)
//  [python]
def test_mag_mask(encoder_list):
    """ Assert identity mask works. """
    for (enc, fb_dim) in encoder_list:
        tf_rep = enc(torch.randn(2, 1, 8000))  # [batch, freq, time]
        id_mag_mask = torch.ones((1, fb_dim // 2, 1))
        masked = transforms.apply_mag_mask(tf_rep, id_mag_mask, dim=1)
        torch.testing.assert_close(masked, tf_rep)
//  [python]
def run_case(bit, B, M, N, C=None, dtype=torch.float32, GS=-1):
    assert bit in [2, 3, 4]
    assert GS == -1 or GS % ({2: 64, 3: 128, 4: 128}[bit]) == 0
    init()
    layer = nn.Linear(M, N)
    vec = torch.randn((B, M) if C is None else (B, C, M)).to("cuda")
    quantizer = Quantizer()

    quantizer.configure(bit=bit, perchannel=True, sym=False, mse=False)
    quantizer.find_params(layer.weight.data, weight=True, groupsize=GS)
    layer.weight.data = quantize(
        layer.weight.data.view(-1, M if GS == -1 else GS),
        quantizer.scale.view(-1, 1),
        quantizer.zero.view(-1, 1),
        quantizer.maxq,
    ).view(N, M)

    qlayer = QuantLinear(layer.in_features, layer.out_features, bit=bit, groupsize=GS)
    qlayer.pack(layer, quantizer.scale, quantizer.zero)

    qlayer = qlayer.to("cuda")
    layer = layer.to("cuda")

    with torch.no_grad():
        gt_out = layer(vec)
        sim_out = qlayer(vec)
        torch.testing.assert_allclose(sim_out, gt_out, rtol=1e-5, atol=1e-5)
//  [python]
def run_case(bit, B, M, N, C=None, dtype=torch.float32, GS=-1):
    assert bit in [2, 3, 4]
    assert GS == -1 or GS % ({2: 64, 3: 128, 4: 128}[bit]) == 0
    init()
    layer = nn.Linear(M, N)
    vec = torch.randn((B, M) if C is None else (B, C, M)).to("cuda")
    quantizer = Quantizer()

    quantizer.configure(bit=bit, perchannel=True, sym=False, mse=False)
    quantizer.find_params(layer.weight.data, weight=True, groupsize=GS)
    layer.weight.data = quantize(
        layer.weight.data.view(-1, M if GS == -1 else GS),
        quantizer.scale.view(-1, 1),
        quantizer.zero.view(-1, 1),
        quantizer.maxq,
    ).view(N, M)

    qlayer = QuantLinear(layer.in_features, layer.out_features, bit=bit, groupsize=GS)
    qlayer.pack(layer, quantizer.scale, quantizer.zero)

    qlayer = qlayer.to("cuda")
    layer = layer.to("cuda")

    with torch.no_grad():
        gt_out = layer(vec)
        sim_out = qlayer(vec)
        torch.testing.assert_close(sim_out, gt_out, rtol=1e-5, atol=1e-5)
//  [python]
def test_set_pred_score(self):
        data_sample = DataSample()
        data_sample.set_pred_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('pred_score', data_sample)
        torch.testing.assert_allclose(data_sample.pred_score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_pred_score(
                torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))
//  [python]
def test_set_pred_score(self):
        data_sample = DataSample()
        data_sample.set_pred_score(torch.tensor([0.1, 0.1, 0.6, 0.1, 0.1]))
        self.assertIn('pred_score', data_sample)
        torch.testing.assert_close(data_sample.pred_score,
                                      [0.1, 0.1, 0.6, 0.1, 0.1])

        # Test invalid length
        with self.assertRaisesRegex(AssertionError, 'should be equal to'):
            data_sample.set_gt_score([1, 2])

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            data_sample.set_pred_score(
                torch.tensor([[0.1, 0.1, 0.6, 0.1, 0.1]]))