            # TODO: is something wrong here? If so, the problem goes deeper
            # n = C1.numel()
            # p = 0.06
            std = C1.std(0).view(1, -1)
            C1 /= std
            C4 /= std
            # assert_all_approx_close(C1, C4, atol=0.02, rtol=0.1, count=int(n*0.06))
            # assert (count / n < p), f"error in more than {p} of elements: {count}/{n}={count/n}"

            C5 = F.int8_mm_dequant(C2, maxA, maxB, bias=bias)
            C5 /= std
            torch.testing.assert_close(C5, C4, atol=0.015, rtol=0.1)
            n = C5.numel()
            assert_all_approx_close(C1, C4, atol=0.015, rtol=0.1, count=int(0.01 * n))

    @pytest.mark.parametrize("dim1", [1 * 1024], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [1 * 1024], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims"))
    @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("decomp"))
    @pytest.mark.deprecated
    def test_colrow_absmax(self, dim1, dim2, dims, threshold):
        for i in range(k):
            A = torch.randn(dim1, dim2, device="cuda").half()

            assert dims == 2

            row_stats1, _ = torch.abs(A.float()).max(1)
            col_stats1, _ = torch.abs(A.float()).max(0)

            if threshold > 0.0:
                A_truncated = A.clone()
                A_truncated[torch.abs(A_truncated) >= threshold] = 0.0
                row_stats1_trunc, _ = torch.abs(A_truncated.float()).max(1)
                col_stats1_trunc, _ = torch.abs(A_truncated.float()).max(0)

                row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=threshold)

                nnz_rows1_counts = (torch.abs(A) >= threshold).sum(1).flatten()
                nnz_block_ptr1 = torch.zeros(
                    nnz_rows1_counts.shape[0] + 1,
                    dtype=nnz_rows1_counts.dtype,
                    device=nnz_rows1_counts.device,
                )
                nnz_block_ptr1[1:] = nnz_rows1_counts.cumsum(0)

                torch.testing.assert_close(col_stats1_trunc, col_stats2)
                torch.testing.assert_close(row_stats1_trunc, row_stats2)
                # torch.testing.assert_close(nnz_block_ptr1, nnz_block_ptr2)
            else:
                row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=0.0)
                assert nnz_block_ptr2 is None
                torch.testing.assert_close(col_stats1, col_stats2)
                torch.testing.assert_close(row_stats1, row_stats2)

    @pytest.mark.parametrize("dim1", [2048, 4096], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [512, 1024], ids=id_formatter("dim2"))
    @pytest.mark.deprecated
    def test_int8_double_quant(self, dim1, dim2):
        for i in range(k):
            A = torch.randn(dim1, dim2, device="cuda").half()
            out_col1, Scol = self.vectorwise_quant(A, dim=0)
            out_row1, Srow = self.vectorwise_quant(A, dim=1)

            CA, CAt, statsA, statsAt, _ = F.int8_double_quant(A)

            # max difference is 1 due to rounding differences
            torch.testing.assert_close(CA, out_row1, atol=1, rtol=0)
            torch.testing.assert_close(CAt, out_col1, atol=1, rtol=0)

            n = CAt.numel()
            num_not_close_rows = (torch.isclose(CA, out_row1, atol=1) == 0).sum().item()
            num_not_close_cols = (torch.isclose(CAt, out_col1, atol=1) == 0).sum().item()


                nnz_rows1_counts = (torch.abs(A) >= threshold).sum(1).flatten()
                nnz_block_ptr1 = torch.zeros(
                    nnz_rows1_counts.shape[0] + 1,
                    dtype=nnz_rows1_counts.dtype,
                    device=nnz_rows1_counts.device,
                )
                nnz_block_ptr1[1:] = nnz_rows1_counts.cumsum(0)

                torch.testing.assert_close(col_stats1_trunc, col_stats2)
                torch.testing.assert_close(row_stats1_trunc, row_stats2)
                # torch.testing.assert_close(nnz_block_ptr1, nnz_block_ptr2)
            else:
                row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=0.0)
                assert nnz_block_ptr2 is None
                torch.testing.assert_close(col_stats1, col_stats2)
                torch.testing.assert_close(row_stats1, row_stats2)

    @pytest.mark.parametrize("dim1", [2048, 4096], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [512, 1024], ids=id_formatter("dim2"))
    @pytest.mark.deprecated
    def test_int8_double_quant(self, dim1, dim2):
        for i in range(k):
            A = torch.randn(dim1, dim2, device="cuda").half()
            out_col1, Scol = self.vectorwise_quant(A, dim=0)
            out_row1, Srow = self.vectorwise_quant(A, dim=1)

            CA, CAt, statsA, statsAt, _ = F.int8_double_quant(A)

            # max difference is 1 due to rounding differences
            torch.testing.assert_close(CA, out_row1, atol=1, rtol=0)
            torch.testing.assert_close(CAt, out_col1, atol=1, rtol=0)

            n = CAt.numel()
            num_not_close_rows = (torch.isclose(CA, out_row1, atol=1) == 0).sum().item()
            num_not_close_cols = (torch.isclose(CAt, out_col1, atol=1) == 0).sum().item()

            # allow for 1:500 error due to rounding differences
            min_error = 1 / 500
            if num_not_close_cols > (min_error * n):
                print(
                    f"Min error exceeded {num_not_close_cols} elements are different. Error: {num_not_close_cols / n:.4f}"
                )
                assert False
            if num_not_close_rows > (min_error * n):
                print(
                    f"Min error exceeded {num_not_close_rows} elements are different. Error: {num_not_close_rows / n:.4f}"
                )
                assert False

            torch.testing.assert_close(Srow.flatten().float(), statsA)
            torch.testing.assert_close(Scol.flatten().float(), statsAt)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize(
        ("dim1", "dim4", "inner"),
        (
            pytest.param(dim1, dim4, inner, id=f"{dim1=},{dim4=},{inner=}")
            for (dim1, dim4, inner) in zip(
                (1, 8, 2048, 4096),
                (2, 128, 2048, 4096),
                (4, 256, 512, 4096),
            )
        ),
    )
    def test_integrated_int8_linear_matmul(self, device, dim1, dim4, inner):
        if device == "cpu" and inner > 2048:
            pytest.skip("Slow on CPU")

        for i in range(k):
            A = torch.randn(dim1, inner, device=device).half()
            B = torch.randn(dim4, inner, device=device).half()


                print(
                    f"Min error exceeded {num_not_close_rows} elements are different. Error: {num_not_close_rows / n:.4f}"
                )
                assert False

            torch.testing.assert_close(Srow.flatten().float(), statsA)
            torch.testing.assert_close(Scol.flatten().float(), statsAt)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize(
        ("dim1", "dim4", "inner"),
        (
            pytest.param(dim1, dim4, inner, id=f"{dim1=},{dim4=},{inner=}")
            for (dim1, dim4, inner) in zip(
                (1, 8, 2048, 4096),
                (2, 128, 2048, 4096),
                (4, 256, 512, 4096),
            )
        ),
    )
    def test_integrated_int8_linear_matmul(self, device, dim1, dim4, inner):
        if device == "cpu" and inner > 2048:
            pytest.skip("Slow on CPU")

        for i in range(k):
            A = torch.randn(dim1, inner, device=device).half()
            B = torch.randn(dim4, inner, device=device).half()

            out1 = torch.matmul(A.half(), B.t().half())

            C1a, stats1a, _ = F.int8_vectorwise_quant(A)
            C2a, stats2a, _ = F.int8_vectorwise_quant(B)
            A1, maxA = self.vectorwise_quant(A, dim=1)
            B1, maxB = self.vectorwise_quant(B, dim=1)

            torch.testing.assert_close(maxA.flatten().float(), stats1a)
            torch.testing.assert_close(maxB.flatten().float(), stats2a)
            torch.testing.assert_close(C1a, A1, rtol=0, atol=1)
            torch.testing.assert_close(C2a, B1, rtol=0, atol=1)

            out2 = F.int8_linear_matmul(A1, B1)

            C2 = F.int8_linear_matmul(A1, B1)

            out3 = self.vectorwise_mm_dequant(C2.float(), maxA, maxB.t())

            err1 = torch.abs(out1 - out2).mean().item()
            err2 = torch.abs(out1 - out3).mean().item()
            assert err2 <= err1 * 1.025

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2"))
    def test_coo_double_quant(self, device, dim1, dim2):
        threshold = 2.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A1 = A * idx
                A2 = torch.zeros_like(A) + A1
                torch.testing.assert_close(A1, A2)

                A[:, outlier_cols] = 0
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                torch.testing.assert_close(A, A2, rtol=0.05, atol=1.5e-2)


            B1, maxB = self.vectorwise_quant(B, dim=1)

            torch.testing.assert_close(maxA.flatten().float(), stats1a)
            torch.testing.assert_close(maxB.flatten().float(), stats2a)
            torch.testing.assert_close(C1a, A1, rtol=0, atol=1)
            torch.testing.assert_close(C2a, B1, rtol=0, atol=1)

            out2 = F.int8_linear_matmul(A1, B1)

            C2 = F.int8_linear_matmul(A1, B1)

            out3 = self.vectorwise_mm_dequant(C2.float(), maxA, maxB.t())

            err1 = torch.abs(out1 - out2).mean().item()
            err2 = torch.abs(out1 - out3).mean().item()
            assert err2 <= err1 * 1.025

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2"))
    def test_coo_double_quant(self, device, dim1, dim2):
        threshold = 2.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A1 = A * idx
                A2 = torch.zeros_like(A) + A1
                torch.testing.assert_close(A1, A2)

                A[:, outlier_cols] = 0
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                torch.testing.assert_close(A, A2, rtol=0.05, atol=1.5e-2)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2"))
    def test_coo_int8_vectorwise_quant(self, device, dim1, dim2):
        threshold = 3.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                A[:, outlier_cols] = 0
                torch.testing.assert_close(A * (idx == 0), A2, rtol=0.05, atol=1.5e-2)

@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")


    def test_coo_double_quant(self, device, dim1, dim2):
        threshold = 2.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A1 = A * idx
                A2 = torch.zeros_like(A) + A1
                torch.testing.assert_close(A1, A2)

                A[:, outlier_cols] = 0
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                torch.testing.assert_close(A, A2, rtol=0.05, atol=1.5e-2)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2"))
    def test_coo_int8_vectorwise_quant(self, device, dim1, dim2):
        threshold = 3.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                A[:, outlier_cols] = 0
                torch.testing.assert_close(A * (idx == 0), A2, rtol=0.05, atol=1.5e-2)

@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
class TestSpMMFunctional:
    @pytest.mark.parametrize("dim1", [256, 1024], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [128, 512], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("transposed_B", TRUE_FALSE, ids=id_formatter("transposed_B"))
    def test_spmm_coo(self, dim1, dim2, transposed_B):
        threshold = 1.5
        dim3 = torch.randint(32, 128, size=(1,)).item()
        # dim3 = 17
        for i in range(k):
            A = torch.randn(dim1, dim2).cuda().half()
            if transposed_B:
                B = torch.randn(dim3, dim2).cuda().half()
            else:
                B = torch.randn(dim2, dim3).cuda().half()

            idx = torch.abs(A) >= threshold


    def test_coo_int8_vectorwise_quant(self, device, dim1, dim2):
        threshold = 3.00
        for i in range(k):
            A = torch.randn(dim1, dim2, device=device).half()

            idx = torch.abs(A) >= threshold
            CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold)

            if outlier_cols is not None:
                A2 = (CA.float() * statsA.unsqueeze(1) / 127).half()
                A[:, outlier_cols] = 0
                torch.testing.assert_close(A * (idx == 0), A2, rtol=0.05, atol=1.5e-2)

@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
class TestSpMMFunctional:
    @pytest.mark.parametrize("dim1", [256, 1024], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [128, 512], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("transposed_B", TRUE_FALSE, ids=id_formatter("transposed_B"))
    def test_spmm_coo(self, dim1, dim2, transposed_B):
        threshold = 1.5
        dim3 = torch.randint(32, 128, size=(1,)).item()
        # dim3 = 17
        for i in range(k):
            A = torch.randn(dim1, dim2).cuda().half()
            if transposed_B:
                B = torch.randn(dim3, dim2).cuda().half()
            else:
                B = torch.randn(dim2, dim3).cuda().half()

            idx = torch.abs(A) >= threshold
            nnz = (idx == 1).sum().item()
            rows, cols = torch.where(idx)
            values = A[idx]
            cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
            A2 = A * idx

            if transposed_B:
                out2 = F.spmm_coo(cooA, B.t())
                out1 = torch.matmul(A2, B.t())
            else:
                out2 = F.spmm_coo(cooA, B)
                out1 = torch.matmul(A2, B)

            assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=30)

    @pytest.mark.benchmark
    def test_spmm_bench(self):
        batch = 2
        model = 1024 * 1
        hidden = model * 4
        seq = 1024
        dim1 = batch * seq
        dim2 = model
        dim3 = hidden
        threshold = 4
        A = torch.randn(dim1, dim2, device="cuda").half()
        B = torch.randn(dim2, dim3, device="cuda").half()
        for i in range(10):
            C1 = bnb.matmul(A, B.t())

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(k):
            C1 = bnb.matmul(A, B.t())


            else:
                B = torch.randn(dim2, dim3).cuda().half()

            idx = torch.abs(A) >= threshold
            nnz = (idx == 1).sum().item()
            rows, cols = torch.where(idx)
            values = A[idx]
            cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
            A2 = A * idx

            if transposed_B:
                out2 = F.spmm_coo(cooA, B.t())
                out1 = torch.matmul(A2, B.t())
            else:
                out2 = F.spmm_coo(cooA, B)
                out1 = torch.matmul(A2, B)

            assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=30)

    @pytest.mark.benchmark
    def test_spmm_bench(self):
        batch = 2
        model = 1024 * 1
        hidden = model * 4
        seq = 1024
        dim1 = batch * seq
        dim2 = model
        dim3 = hidden
        threshold = 4
        A = torch.randn(dim1, dim2, device="cuda").half()
        B = torch.randn(dim2, dim3, device="cuda").half()
        for i in range(10):
            C1 = bnb.matmul(A, B.t())

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(k):
            C1 = bnb.matmul(A, B.t())
        torch.cuda.synchronize()
        t8 = time.time() - t0

        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        print(nnz / idx.numel())
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)

        for i in range(10):
            out2 = F.spmm_coo(cooA, B)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(k):
            out2 = F.spmm_coo(cooA, B)
        torch.cuda.synchronize()
        tsp = time.time() - t0
        print(tsp, t8)
        print(tsp / t8)

    @pytest.mark.parametrize("dim1", [1 * 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [12288], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("dtype", [torch.float16], ids=describe_dtype)
    @pytest.mark.parametrize("out_func", ["zeros", "ones"], ids=id_formatter("out_func"))
    def test_spmm_coo_very_sparse(self, dim1, dim2, dtype, out_func):
        out_func = getattr(torch, out_func)

        threshold = 3.3
        # threshold = 2.8
        # threshold = 0.0
        A = torch.randn(dim1, dim2, device="cuda").half()
        if dtype == torch.float16:
            B = torch.randn(dim2, dim2 * 4, device="cuda").half()
            torch.nn.init.xavier_uniform_(B)
        else:
            B = torch.randn(dim2, dim2 * 4, device="cuda").half()
            torch.nn.init.xavier_uniform_(B)

            SB = torch.abs(B).max().float()


        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)

        for i in range(10):
            out2 = F.spmm_coo(cooA, B)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(k):
            out2 = F.spmm_coo(cooA, B)
        torch.cuda.synchronize()
        tsp = time.time() - t0
        print(tsp, t8)
        print(tsp / t8)

    @pytest.mark.parametrize("dim1", [1 * 2048], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [12288], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("dtype", [torch.float16], ids=describe_dtype)
    @pytest.mark.parametrize("out_func", ["zeros", "ones"], ids=id_formatter("out_func"))
    def test_spmm_coo_very_sparse(self, dim1, dim2, dtype, out_func):
        out_func = getattr(torch, out_func)

        threshold = 3.3
        # threshold = 2.8
        # threshold = 0.0
        A = torch.randn(dim1, dim2, device="cuda").half()
        if dtype == torch.float16:
            B = torch.randn(dim2, dim2 * 4, device="cuda").half()
            torch.nn.init.xavier_uniform_(B)
        else:
            B = torch.randn(dim2, dim2 * 4, device="cuda").half()
            torch.nn.init.xavier_uniform_(B)

            SB = torch.abs(B).max().float()
            B = torch.round(B / SB * 127).to(torch.int8)

        print("")
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        out1 = torch.matmul(A2.half(), B.half())
        out = out_func(out1.shape, dtype=torch.float16, device=out1.device)
        out1 += out.clone()
        out2 = F.spmm_coo_very_sparse(cooA, B, out=out)
        # print(B)
        # print(out1)
        # print(out2)
        p = 200 / (2048 * 12288 * 4)
        n = out1.numel()
        count = math.ceil(p * n)
        std = out1.std()
        out1 /= std
        out2 /= std
        assert_all_approx_close(out1, out2.half(), rtol=0.01, atol=3.0e-2, count=count)
        # assert_all_approx_close(out1, out2.half(), rtol=0.05, atol=0.01, count=count)

        idx_col = torch.randint(0, A2.shape[-1], size=(15,))

        # torch.testing.assert_close(out1, out2.half(), rtol=0.05, atol=0.001)

        # Bt = torch.randn(dim2*4, dim2, device='cuda').half()
        # torch.cuda.synchronize()
        # t0 = time.time()
        # print(A2.shape, B.shape)
        # for i in range(100):
        #   #out3 = F.spmm_coo(cooA, Bt.t())
        #   #out2 = F.spmm_coo(cooA, B)
        #   #out2 = F.spmm_coo_very_sparse(cooA, B)
        #   #out1 = torch.matmul(A, Bt.t())

        # torch.cuda.synchronize()
        # print(time.time() - t0)

    @pytest.mark.parametrize("dim1", [1 * 2048])
    @pytest.mark.parametrize("dim2", [2048])
    @pytest.mark.parametrize("dtype", [torch.int8])
    def test_spmm_coo_dequant(self, dim1, dim2, dtype):


        idx_col = torch.randint(0, A2.shape[-1], size=(15,))

        # torch.testing.assert_close(out1, out2.half(), rtol=0.05, atol=0.001)

        # Bt = torch.randn(dim2*4, dim2, device='cuda').half()
        # torch.cuda.synchronize()
        # t0 = time.time()
        # print(A2.shape, B.shape)
        # for i in range(100):
        #   #out3 = F.spmm_coo(cooA, Bt.t())
        #   #out2 = F.spmm_coo(cooA, B)
        #   #out2 = F.spmm_coo_very_sparse(cooA, B)
        #   #out1 = torch.matmul(A, Bt.t())

        # torch.cuda.synchronize()
        # print(time.time() - t0)

    @pytest.mark.parametrize("dim1", [1 * 2048])
    @pytest.mark.parametrize("dim2", [2048])
    @pytest.mark.parametrize("dtype", [torch.int8])
    def test_spmm_coo_dequant(self, dim1, dim2, dtype):
        threshold = 6.0
        # threshold = 2.8
        # threshold = 0.0
        A = torch.randn(dim1, dim2, device="cuda").half()
        B = torch.empty(dim2, dim2 * 4, device="cuda", dtype=torch.float16)
        torch.nn.init.xavier_uniform_(B)
        Bt = B.t().contiguous()

        CB, CBt, statsB, statsBt, coo_tensor = F.int8_double_quant(B)

        rowidx = torch.randint(0, A.shape[-1], size=(15,))

        A[:, rowidx] = 8.0

        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
        out1 = torch.matmul(A2, B.half())
        out3 = F.spmm_coo_very_sparse(cooA, CBt.half())
        out3 = out3 * statsBt.half() / 127

        values, counts = torch.unique(cooA.rowidx, return_counts=True)
        offset = counts.cumsum(0).int()
        max_count, max_idx = torch.sort(counts, descending=True)
        print(torch.median(max_count.float()))

        torch.testing.assert_close(out2, out3, rtol=0.05, atol=0.001)

        p = 200 / (2048 * 12288 * 4)
        n = out1.numel()
        count = math.ceil(p * n)
        assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=count)

        # torch.cuda.synchronize()
        # t0 = time.time()
        # for i in range(100):
        #   out2 = F.spmm_coo_very_sparse(cooA, B)
        # torch.cuda.synchronize()
        # print('fp16', time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out2 = F.spmm_coo(cooA, B)
        torch.cuda.synchronize()
        print("cusparse fp16", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out2 = F.spmm_coo_very_sparse(cooA, CBt)
        torch.cuda.synchronize()
        print("int8", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
        torch.cuda.synchronize()
        print("int8+dequant", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out2 = torch.matmul(A, B)
        torch.cuda.synchronize()
        print("matmul", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out1 = bnb.matmul(A, Bt)
            out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
            out = out1 + out2
        torch.cuda.synchronize()
        print("sparse+ matmul", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out1 = bnb.matmul(A, Bt)
            torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1)
        torch.cuda.synchronize()
        print("partial matmul", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out1 = bnb.matmul(A, Bt)
        torch.cuda.synchronize()
        print("partial matmul", time.time() - t0)

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
class TestSparseTensorFunctional:
    def test_coo2csr(self):
        threshold = 1
        A = torch.randn(128, 128).half().cuda()
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        csrA = F.coo2csr(cooA)
        counts = csrA.rowptr[1:] - csrA.rowptr[:-1]
        assert counts.numel() == A.shape[0]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(1))
        idx = A2 != 0
        torch.testing.assert_close(A2[idx], csrA.values)


        print("sparse+ matmul", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out1 = bnb.matmul(A, Bt)
            torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1)
        torch.cuda.synchronize()
        print("partial matmul", time.time() - t0)

        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(100):
            out1 = bnb.matmul(A, Bt)
        torch.cuda.synchronize()
        print("partial matmul", time.time() - t0)

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
class TestSparseTensorFunctional:
    def test_coo2csr(self):
        threshold = 1
        A = torch.randn(128, 128).half().cuda()
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        csrA = F.coo2csr(cooA)
        counts = csrA.rowptr[1:] - csrA.rowptr[:-1]
        assert counts.numel() == A.shape[0]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(1))
        idx = A2 != 0
        torch.testing.assert_close(A2[idx], csrA.values)

    def test_coo2csc(self):
        threshold = 1
        A = torch.randn(128, 128).half().cuda()
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        cscA = F.coo2csc(cooA)
        counts = cscA.colptr[1:] - cscA.colptr[:-1]
        assert counts.numel() == A.shape[1]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(0))
        # torch uses row-major -> use transpose to transfer to col-major
        idx = A2.t() != 0
        torch.testing.assert_close(A2.t()[idx], cscA.values)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
class TestSparseTensorFunctional:
    def test_coo2csr(self):
        threshold = 1
        A = torch.randn(128, 128).half().cuda()
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        csrA = F.coo2csr(cooA)
        counts = csrA.rowptr[1:] - csrA.rowptr[:-1]
        assert counts.numel() == A.shape[0]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(1))
        idx = A2 != 0
        torch.testing.assert_close(A2[idx], csrA.values)

    def test_coo2csc(self):
        threshold = 1
        A = torch.randn(128, 128).half().cuda()
        idx = torch.abs(A) >= threshold
        nnz = (idx == 1).sum().item()
        rows, cols = torch.where(idx)
        values = A[idx]
        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        cscA = F.coo2csc(cooA)
        counts = cscA.colptr[1:] - cscA.colptr[:-1]
        assert counts.numel() == A.shape[1]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(0))
        # torch uses row-major -> use transpose to transfer to col-major
        idx = A2.t() != 0
        torch.testing.assert_close(A2.t()[idx], cscA.values)

class TestQuantize4BitFunctional:
    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype)
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize(
        "blocksize",
        [64, 128, 256, 512, 1024, 2048, 4096] if not HIP_ENVIRONMENT else [128, 256, 512, 1024, 2048, 4096],
    )
    def test_4bit_quant(self, device, dtype, quant_type, blocksize):
        if device == "hpu" and not is_supported_on_hpu(quant_type, dtype):
            pytest.skip("This configuration is not supported on HPU.")

        A1 = torch.randn(1024, 1024, device=device, dtype=dtype)
        qa, SA = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type)
        A2 = F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type)

        err = (A1 - A2).abs().float()
        relerr = (err / (A1.abs().float() + 1e-8)).mean()


        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
        A2 = A * idx
        cscA = F.coo2csc(cooA)
        counts = cscA.colptr[1:] - cscA.colptr[:-1]
        assert counts.numel() == A.shape[1]

        torch.testing.assert_close(counts.long(), (A2 != 0).sum(0))
        # torch uses row-major -> use transpose to transfer to col-major
        idx = A2.t() != 0
        torch.testing.assert_close(A2.t()[idx], cscA.values)

class TestQuantize4BitFunctional:
    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype)
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize(
        "blocksize",
        [64, 128, 256, 512, 1024, 2048, 4096] if not HIP_ENVIRONMENT else [128, 256, 512, 1024, 2048, 4096],
    )
    def test_4bit_quant(self, device, dtype, quant_type, blocksize):
        if device == "hpu" and not is_supported_on_hpu(quant_type, dtype):
            pytest.skip("This configuration is not supported on HPU.")

        A1 = torch.randn(1024, 1024, device=device, dtype=dtype)
        qa, SA = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type)
        A2 = F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type)

        err = (A1 - A2).abs().float()
        relerr = (err / (A1.abs().float() + 1e-8)).mean()
        err = err.mean()

        assert A2.dtype == dtype

        # With larger block sizes, we can expect this to blow up.
        # At blocksize>=1024, don't even bother looking at relerr.
        #
        # Actually, the above is not true anymore after fixing the integer packing bug.
        # The following values were taken from averaging 1k samples per test configuration after fixing the bug.
        error_dict = dict()
        error_dict["fp4"] = dict()
        error_dict["nf4"] = dict()
        error_dict["fp4"]["err"] = {
            64: 0.096545,
            128: 0.102947,
            256: 0.108685,
            512: 0.114087,
            1024: 0.119312,
            2048: 0.124460,
            4096: 0.129573,
        }
        error_dict["fp4"]["rel_err"] = {
            64: 0.260130,
            128: 0.275734,
            256: 0.289842,
            512: 0.302852,
            1024: 0.314982,
            2048: 0.326402,
            4096: 0.337228,
        }

        error_dict["nf4"]["err"] = {
            64: 0.072792,
            128: 0.076835,
            256: 0.080326,
            512: 0.083535,
            1024: 0.086603,
            2048: 0.089592,
            4096: 0.092537,
        }
        error_dict["nf4"]["rel_err"] = {
            64: 0.203299,
            128: 0.215252,
            256: 0.226044,
            512: 0.236021,
            1024: 0.245365,
            2048: 0.254146,
            4096: 0.262457,
        }

        # Allow higher tolerance for fp32 on CPU with larger block sizes
        reltol = 2.8e-3 if dtype == torch.float32 and blocksize >= 128 and device == "cpu" else 1e-3
        errtol = 1.2e-3 if dtype == torch.float32 and blocksize >= 1024 and device == "cpu" else 1e-3

        assert err < error_dict[quant_type]["err"][blocksize] + errtol
        assert relerr < error_dict[quant_type]["rel_err"][blocksize] + reltol

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128], ids=id_formatter("blocksize"))
    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=describe_dtype)
    def test_4bit_compressed_stats(self, device, quant_type, blocksize, dtype):
        if device == "hpu" and not is_supported_on_hpu(quant_type, dtype):
            pytest.skip("FP4 quantization is not supported on HPU.")

        errs1 = []
        errs2 = []
        for i in range(10):
            A1 = torch.randn(1024, 1024, device=device, dtype=dtype)
            q2, SA2 = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type)
            q3, SA3 = F.quantize_4bit(A1, blocksize=blocksize, compress_statistics=True, quant_type=quant_type)
            A2 = F.dequantize_4bit(q2, SA2, quant_type=quant_type)
            A3 = F.dequantize_4bit(q3, SA3, quant_type=quant_type)

            err = (A1 - A2).abs().float()
            relerr = (err / (A1.abs().float() + 1e-15)).mean()


            64: 0.203299,
            128: 0.215252,
            256: 0.226044,
            512: 0.236021,
            1024: 0.245365,
            2048: 0.254146,
            4096: 0.262457,
        }

        # Allow higher tolerance for fp32 on CPU with larger block sizes
        reltol = 2.8e-3 if dtype == torch.float32 and blocksize >= 128 and device == "cpu" else 1e-3
        errtol = 1.2e-3 if dtype == torch.float32 and blocksize >= 1024 and device == "cpu" else 1e-3

        assert err < error_dict[quant_type]["err"][blocksize] + errtol
        assert relerr < error_dict[quant_type]["rel_err"][blocksize] + reltol

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128], ids=id_formatter("blocksize"))
    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=describe_dtype)
    def test_4bit_compressed_stats(self, device, quant_type, blocksize, dtype):
        if device == "hpu" and not is_supported_on_hpu(quant_type, dtype):
            pytest.skip("FP4 quantization is not supported on HPU.")

        errs1 = []
        errs2 = []
        for i in range(10):
            A1 = torch.randn(1024, 1024, device=device, dtype=dtype)
            q2, SA2 = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type)
            q3, SA3 = F.quantize_4bit(A1, blocksize=blocksize, compress_statistics=True, quant_type=quant_type)
            A2 = F.dequantize_4bit(q2, SA2, quant_type=quant_type)
            A3 = F.dequantize_4bit(q3, SA3, quant_type=quant_type)

            err = (A1 - A2).abs().float()
            relerr = (err / (A1.abs().float() + 1e-15)).mean()
            err = err.mean()

            errs1.append(err.item())

            assert err.item() < 0.11
            assert relerr.item() < 0.28

            err = (A1 - A3).abs().float()
            relerr = (err / (A1.abs().float() + 1e-15)).mean()
            err = err.mean()

            errs2.append(err.item())

            assert err.item() < 0.11
            assert relerr.item() < 0.28

    # @pytest.mark.parametrize("quant_type", ['fp4', 'nf4'])
    @pytest.mark.parametrize("quant_type", ["nf4"])
    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
    @pytest.mark.benchmark
    def test_bench_4bit_dequant(self, quant_type):
        blocksize = 256
        a = torch.rand(1024 * 12 * 4, 1024 * 12, device="cuda").half()
        qa, SA = F.quantize_4bit(a, blocksize=blocksize, quant_type=quant_type)

        input_size = a.numel() / 2
        output_size = a.numel() * 2
        num_bytes = input_size + output_size
        GB = num_bytes / 1e9
        max_theoretical_s = GB / 768
        # print(max_theoretical_s*1e6)
        b = torch.randn(128, 1024 * 12, device="cuda").half()

        iters = 100
        torch.cuda.synchronize()


            err = err.mean()

            errs1.append(err.item())

            assert err.item() < 0.11
            assert relerr.item() < 0.28

            err = (A1 - A3).abs().float()
            relerr = (err / (A1.abs().float() + 1e-15)).mean()
            err = err.mean()

            errs2.append(err.item())

            assert err.item() < 0.11
            assert relerr.item() < 0.28

    # @pytest.mark.parametrize("quant_type", ['fp4', 'nf4'])
    @pytest.mark.parametrize("quant_type", ["nf4"])
    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
    @pytest.mark.benchmark
    def test_bench_4bit_dequant(self, quant_type):
        blocksize = 256
        a = torch.rand(1024 * 12 * 4, 1024 * 12, device="cuda").half()
        qa, SA = F.quantize_4bit(a, blocksize=blocksize, quant_type=quant_type)

        input_size = a.numel() / 2
        output_size = a.numel() * 2
        num_bytes = input_size + output_size
        GB = num_bytes / 1e9
        max_theoretical_s = GB / 768
        # print(max_theoretical_s*1e6)
        b = torch.randn(128, 1024 * 12, device="cuda").half()

        iters = 100
        torch.cuda.synchronize()
        t0 = time.time()
        for i in range(iters):
            F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type)
            # b.copy_(a)
        torch.cuda.synchronize()
        # print((time.time()-t0)/iters*1e6)

        # torch.cuda.synchronize()
        # t0 = time.time()
        # for i in range(iters):
        #    torch.matmul(b, a.t())
        # torch.cuda.synchronize()
        # print((time.time()-t0)/iters*1e6)

    @pytest.mark.skipif(
        HIP_ENVIRONMENT, reason="gemv 4bit tests are partially enabled on MI300, others being fixed for warpsize 64"
    )
    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("double_quant", TRUE_FALSE, ids=lambda double_quant: f"DQ_{double_quant}")
    @pytest.mark.parametrize("storage_type", ["nf4", "fp4"])
    @pytest.mark.parametrize("kind", ["fc1", "fc2", "attn", "attn_packed"])
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
    @pytest.mark.parametrize(
        "quant_storage",
        [torch.uint8, torch.float16, torch.bfloat16, torch.float32],


        # t0 = time.time()
        # for i in range(iters):
        #    torch.matmul(b, a.t())
        # torch.cuda.synchronize()
        # print((time.time()-t0)/iters*1e6)

    @pytest.mark.skipif(
        HIP_ENVIRONMENT, reason="gemv 4bit tests are partially enabled on MI300, others being fixed for warpsize 64"
    )
    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("double_quant", TRUE_FALSE, ids=lambda double_quant: f"DQ_{double_quant}")
    @pytest.mark.parametrize("storage_type", ["nf4", "fp4"])
    @pytest.mark.parametrize("kind", ["fc1", "fc2", "attn", "attn_packed"])
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
    @pytest.mark.parametrize(
        "quant_storage",
        [torch.uint8, torch.float16, torch.bfloat16, torch.float32],
        ids=describe_dtype,
    )
    @pytest.mark.parametrize("dim", [128, 256, 512, 1024], ids=id_formatter("dim"))
    def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double_quant, kind):
        if device == "hpu" and not is_supported_on_hpu(storage_type, dtype, quant_storage):
            pytest.skip("This configuration is not supported on HPU.")

        errs1 = []
        errs2 = []
        errs3 = []
        relerrs1 = []
        relerrs2 = []
        relerrs3 = []
        max_errs1 = []
        max_errs2 = []
        max_errs3 = []

        # Large number of iterations is excessive and slow on CPU.
        # Keep for CUDA/XPU for now.
        iters = 10 if device == "cpu" else 100

        for i in range(iters):
            if kind == "fc1":
                A = torch.randn(1, dim, dtype=dtype, device=device)
                B = torch.randn(dim * 4, dim, dtype=dtype, device=device) / math.sqrt(dim)
            elif kind == "fc2":
                A = torch.randn(1, 4 * dim, dtype=dtype, device=device)
                B = torch.randn(dim, 4 * dim, dtype=dtype, device=device) / math.sqrt(dim)
            elif kind == "attn":
                A = torch.randn(1, dim, dtype=dtype, device=device)
                B = torch.randn(dim, dim, dtype=dtype, device=device) / math.sqrt(dim)
            elif kind == "attn_packed":
                A = torch.randn(1, dim, dtype=dtype, device=device)
                B = torch.randn(dim * 3, dim, dtype=dtype, device=device) / math.sqrt(dim)

            qB, state = F.quantize_4bit(
                B,
                quant_type=storage_type,
                compress_statistics=double_quant,
                quant_storage=quant_storage,
            )
            C3 = torch.matmul(A, B.t())
            C2 = F.gemv_4bit(A, qB.t(), state=state)
            A.requires_grad = True
            C1 = bnb.matmul_4bit(A, qB.t(), state)

            err1 = (C1 - C2).abs().float()
            err2 = (C3 - C2).abs().float()
            err3 = (C3 - C1).abs().float()

            mag1 = torch.abs(C1).float() + 1e-5
            mag2 = torch.abs(C3).float() + 1e-5
            mag3 = torch.abs(C3).float() + 1e-5

            relerr1 = err1 / mag1
            relerr2 = err2 / mag2
            relerr3 = err3 / mag3

            max_err1 = err1.max()
            max_err2 = err2.max()
            max_err3 = err3.max()

            errs1.append(err1.mean().item())
            errs2.append(err2.mean().item())
            errs3.append(err3.mean().item())

            relerrs1.append(relerr1.mean().item())
            relerrs2.append(relerr2.mean().item())
            relerrs3.append(relerr3.mean().item())

            max_errs1.append(max_err1.item())
            max_errs2.append(max_err2.item())
            max_errs3.append(max_err3.item())

            c = int(C1.numel() * 0.0014 * (dim / 256)) + 1

            c = assert_all_approx_close(C1, C2, 1e-5, 0.01, count=0, throw=False)
        err1 = sum(errs1) / len(errs1) / math.sqrt(dim)
        err2 = sum(errs2) / len(errs2) / math.sqrt(dim)
        err3 = sum(errs3) / len(errs3) / math.sqrt(dim)
        relerr1 = sum(relerrs1) / len(relerrs1) / math.sqrt(dim)
        relerr2 = sum(relerrs2) / len(relerrs2) / math.sqrt(dim)
        relerr3 = sum(relerrs3) / len(relerrs3) / math.sqrt(dim)
        maxerr1 = sum(max_errs1) / len(max_errs1) / math.sqrt(dim)
        maxerr2 = sum(max_errs2) / len(max_errs2) / math.sqrt(dim)
        maxerr3 = sum(max_errs3) / len(max_errs3) / math.sqrt(dim)
        absratio = err2 / err3
        relratio = relerr2 / relerr3
        maxratio = relerr2 / relerr3

        # for debugging if the tests fails
        #
        # print('='*80)
        # print(f'For matmul: {A.shape}, {B.shape}, {kind}, {dtype}, {storage_type}, double_quant={double_quant}:')
        # print(C1.flatten()[-20:])
        # print(C2.flatten()[-20:])
        # print(f'inference vs training abs: {err1}')
        # print(f'inference vs training rel: {relerr1}')
        # print(f'inference vs training max: {maxerr1}')
        # print(f'inference vs training vs torch err ratio abs: {absratio}')
        # print(f'inference vs training vs torch err ratio rel: {relratio}')
        # print(f'inference vs training vs torch err ratio max: {maxratio}')
        if dtype == torch.float16:
            if dim <= 512:
                assert err1 < 7e-5

                # TODO(matthewdouglas): On T4, dim=128-fp16-fc2-fp4-DQ will have relerror ~ 0.00092727
                if (
                    device == "cuda"
                    and double_quant
                    and storage_type == "fp4"
                    and kind == "fc2"
                    and torch.cuda.get_device_capability() == (7, 5)
                ):
                    assert relerr1 < 0.00093
                else:
                    assert relerr1 < 0.0008
            else:
                assert err1 < 6e-5
                assert relerr1 < 2e-4
            assert absratio < 1.005 and absratio > 0.995
            assert relratio < 1.005 and relratio > 0.992
            assert maxratio < 1.005 and maxratio > 0.992
        elif dtype == torch.float32:
            if dim <= 512:
                assert err1 < 5e-8
                assert relerr1 < 1e-6
                assert maxerr1 < 1.05e-7
            else:
                assert err1 < 5e-8
                assert relerr1 < 8e-6
                assert maxerr1 < 1e-7
            assert absratio < 1.005 and absratio > 0.995
            assert relratio < 1.005 and relratio > 0.995
            assert maxratio < 1.005 and maxratio > 0.995
        elif dtype == torch.bfloat16:
            if dim <= 512:
                relerr_thres = 0.013 if hasattr(torch, "xpu") and torch.xpu.is_available() else 0.007
                assert err1 < 6e-4
                assert relerr1 < relerr_thres
                assert maxerr1 < 0.015
            else:
                assert err1 < 2e-4
                assert relerr1 < 0.002
                assert maxerr1 < 0.0012
            assert absratio < 1.005 and absratio > 0.995
            assert relratio < 1.05 and relratio > 0.96
            assert maxratio < 1.05 and maxratio > 0.97

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("storage_type", ["nf4", "fp4"], ids=["nf4", "fp4"])
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
    @pytest.mark.skipif(
        HIP_ENVIRONMENT and ROCM_GPU_ARCH == "gfx90a",
        reason="this test is not supported on ROCm with gfx90a architecture yet",
    )
    def test_gemv_eye_4bit(self, device, storage_type, dtype):
        if device == "hpu" and not is_supported_on_hpu(storage_type, dtype):
            pytest.skip("This configuration is not supported on HPU.")

        if (
            device == "cpu"
            and platform.system() == "Windows"
            and version.parse(torch.__version__).release == (2, 8, 0)
        ):
            pytest.skip("Regression: CPU crash on Windows with torch 2.8.0")

        dims = 4


            if dim <= 512:
                relerr_thres = 0.013 if hasattr(torch, "xpu") and torch.xpu.is_available() else 0.007
                assert err1 < 6e-4
                assert relerr1 < relerr_thres
                assert maxerr1 < 0.015
            else:
                assert err1 < 2e-4
                assert relerr1 < 0.002
                assert maxerr1 < 0.0012
            assert absratio < 1.005 and absratio > 0.995
            assert relratio < 1.05 and relratio > 0.96
            assert maxratio < 1.05 and maxratio > 0.97

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("storage_type", ["nf4", "fp4"], ids=["nf4", "fp4"])
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
    @pytest.mark.skipif(
        HIP_ENVIRONMENT and ROCM_GPU_ARCH == "gfx90a",
        reason="this test is not supported on ROCm with gfx90a architecture yet",
    )
    def test_gemv_eye_4bit(self, device, storage_type, dtype):
        if device == "hpu" and not is_supported_on_hpu(storage_type, dtype):
            pytest.skip("This configuration is not supported on HPU.")

        if (
            device == "cpu"
            and platform.system() == "Windows"
            and version.parse(torch.__version__).release == (2, 8, 0)
        ):
            pytest.skip("Regression: CPU crash on Windows with torch 2.8.0")

        dims = 4
        dims = get_test_dims(0, 8192, n=dims)
        dims = [dim + (64 - (dim % 64)) for dim in dims]
        # for dim in [576, 5120, 3520, 5184, 1280, 4992, 5312, 2048]:
        for dim in dims:
            A = torch.normal(0, 0.1, size=(1, 1, dim), dtype=dtype, device=device)
            B = torch.eye(dim, dtype=dtype, device=device)

            qB, state = F.quantize_4bit(B, quant_type=storage_type, compress_statistics=False)
            C3 = torch.matmul(A, B.t())
            C2 = bnb.matmul_4bit(A, qB.t(), state)
            A.requires_grad = True
            C1 = bnb.matmul_4bit(A, qB.t(), state)

            torch.testing.assert_close(A, C3)
            torch.testing.assert_close(A, C1)
            torch.testing.assert_close(A, C2)
        # torch.testing.assert_close(A, C1, rtol=1e-5, atol=0.00001)
        # torch.testing.assert_close(A, C2, rtol=1e-5, atol=0.080)

def test_normal_map_tree():
    code = F.create_normal_map()
    values = code[:8].tolist() + code[-8:].tolist()
    num_pivots = 1
    # print(values)
    while num_pivots < 16:
        idx = list(range(16 // num_pivots // 2, 16, 16 // num_pivots))
        # print(idx)
        num_pivots *= 2
        pivots = []
        for i in idx:
            pivots.append((values[i - 1] + values[i]) / 2)
        # print(pivots)


    @pytest.mark.parametrize("hidden_dim", [64, 1024, 4096], ids=id_formatter("hidden_dim"))
    @pytest.mark.parametrize("batch_dim", [2, 8, 16], ids=id_formatter("batch_dim"))
    def test_dim3_igemm(self, seq_dim, hidden_dim, batch_dim):
        seq_dim = seq_dim - (seq_dim % 32)
        hidden_dim = hidden_dim - (hidden_dim % 32)
        batch_dim = batch_dim - (batch_dim % 2)
        for i in range(25):
            A = torch.randint(-128, 127, size=(batch_dim, seq_dim, hidden_dim), device="cuda").to(torch.int8)
            B = torch.randint(-128, 127, size=(batch_dim, seq_dim, 1024), device="cuda").to(torch.int8)
            out2 = torch.einsum("bsi, bso->io", A.float(), B.float())
            iout = torch.empty(A.shape[2], B.shape[2], dtype=torch.int32, device=A.device)
            out = F.igemm(A, B, out=iout)

            torch.testing.assert_close(out.float(), out2)

    @pytest.mark.parametrize("seq_dim", [32, 512], ids=id_formatter("seq_dim"))
    @pytest.mark.parametrize("hidden_dim", [32, 1024 * 4], ids=id_formatter("hidden_dim"))
    @pytest.mark.parametrize("batch_dim", [2, 16], ids=id_formatter("batch_dim"))
    @pytest.mark.parametrize("transpose", TRUE_FALSE, ids=id_formatter("transpose"))
    def test_minmax_igemm(self, seq_dim, hidden_dim, batch_dim, transpose):
        def min_max(x):
            maxA = torch.amax(x, dim=2, keepdim=True)
            minA = torch.amin(x, dim=2, keepdim=True)
            scale = (maxA - minA) / 2.0
            return (127 * (x - minA - scale) / scale).to(torch.int8), minA, scale

        seq_dim = seq_dim - (seq_dim % 16)
        hidden_dim = hidden_dim - (hidden_dim % 16)
        batch_dim = batch_dim - (batch_dim % 2)
        errs = []
        relerrs = []
        errs2 = []
        relerrs2 = []
        for i in range(k):
            A = torch.normal(0.0, 0.5, size=(batch_dim, seq_dim, hidden_dim), device="cuda")
            if transpose:
                B = torch.normal(0, 0.5, size=(256, hidden_dim), device="cuda")
            else:
                B = torch.normal(0, 0.5, size=(hidden_dim, 256), device="cuda")
            Ac, minA, scale = min_max(A)
            if transpose:
                maxB, Bc = quant_multi(B, dim=(1 if transpose else 0))
                out = F.igemm(Ac, Bc.t())
                out2 = torch.matmul(A, B.t())
                offset = B.t().sum(0) * (minA + scale)


import copy
import os
import pickle
import platform
from tempfile import TemporaryDirectory

import pytest
import torch

import bitsandbytes as bnb
from bitsandbytes.cextension import HIP_ENVIRONMENT
from tests.helpers import (
    TRUE_FALSE,
    describe_dtype,
    get_available_devices,
    id_formatter,
    is_supported_on_hpu,
    torch_load_from_buffer,
    torch_save_to_buffer,
)

storage = {
    "uint8": torch.uint8,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_storage", ["uint8", "float16", "bfloat16", "float32"])
@pytest.mark.parametrize("original_dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward"))
def test_linear_serialization(
    device, quant_type, original_dtype, compress_statistics, bias, quant_storage, save_before_forward
):
    if device == "hpu" and not is_supported_on_hpu(quant_type, original_dtype, storage[quant_storage]):
        pytest.skip("This configuration is not supported on HPU.")

    compute_dtype = None
    layer_shape = (300, 400)

    linear = torch.nn.Linear(*layer_shape, dtype=original_dtype, device="cpu")  # original layer

    # Quantizing original layer
    linear_q = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        device="meta",
    )
    new_weight = bnb.nn.Params4bit(data=linear.weight, quant_type=quant_type, requires_grad=False)
    linear_q.weight = new_weight
    if bias:
        linear_q.bias = torch.nn.Parameter(linear.bias)
    linear_q = linear_q.to(device)

    # saving to state_dict:
    sd = linear_q.state_dict()

    # restoring from state_dict:
    bias_data2 = sd.pop("bias", None)
    weight_data2 = sd.pop("weight")
    weight2 = bnb.nn.Params4bit.from_prequantized(quantized_stats=sd, data=weight_data2, device=device)

    # creating new layer with same params:
    linear_q2 = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        device="meta",
    )
    # loading weights from state_dict:
    linear_q2.weight = weight2
    if bias:
        linear_q2.bias = torch.nn.Parameter(bias_data2)
    linear_q2 = linear_q2.to(device)

    # MATCHING
    a, b = linear_q.weight, linear_q2.weight

    # Quantizing original layer with specified quant_storage type
    linear_qs = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        quant_storage=storage[quant_storage],
        device="meta",
    )
    linear_qs.weight = bnb.nn.Params4bit(
        data=linear.weight,
        requires_grad=False,
        quant_type=quant_type,
        quant_storage=storage[quant_storage],
    )
    if bias:
        linear_qs.bias = torch.nn.Parameter(linear.bias)
    linear_qs = linear_qs.to(device)

    assert a.device == b.device
    assert a.dtype == b.dtype
    assert torch.equal(a, b)

    q0 = a.quant_state
    q1 = b.quant_state
    for attr in ("code", "dtype", "blocksize", "absmax"):
        c, d = getattr(q0, attr), getattr(q1, attr)
        if isinstance(c, torch.Tensor):
            assert torch.equal(c, d)
        else:
            assert c == d, f"{c} != {d}"

    if q0.state2 is not None:
        for attr in ("code", "dtype", "blocksize", "absmax"):
            c, d = getattr(q0.state2, attr), getattr(q1.state2, attr)
            if isinstance(c, torch.Tensor):
                assert torch.equal(c, d)
            else:
                assert c == d, f"{c} != {d}"

    if bias:
        a, b = linear_q.bias, linear_q2.bias
        assert a.device == b.device
        assert a.dtype == b.dtype
        assert torch.equal(a, b)

    if save_before_forward:
        bytes_4bit = torch_save_to_buffer(linear_q)

    # Forward test
    x = torch.rand(42, layer_shape[0], device=device)
    a = linear_q(x)
    b = linear_q2(x)
    c = linear_qs(x)
    assert a.device == b.device
    assert a.dtype == b.dtype
    assert a.device == c.device
    assert a.dtype == c.dtype
    assert torch.equal(a, b)
    assert torch.equal(a, c)

    if not save_before_forward:
        bytes_4bit = torch_save_to_buffer(linear_q)
    linear_q3 = torch_load_from_buffer(bytes_4bit)

    # Test moving to CPU and back to GPU
    if device != "cpu":
        linear_q2.to("cpu")
        linear_q2.to(device)
    d = linear_qs(x)
    assert c.dtype == d.dtype
    assert c.device == d.device
    assert torch.equal(c, d)

    d = linear_q3(x)
    assert c.dtype == d.dtype
    assert c.device == d.device
    assert torch.equal(c, d)

    # Saved size ratio test. Target set for layer_shape == (300, 400) w/ bias
    with TemporaryDirectory() as tmpdir:
        state_path_4bit = os.path.join(tmpdir, "state_4bit.pth")
        state_path = os.path.join(tmpdir, "state.pth")
        torch.save(linear.state_dict(), state_path)
        torch.save(linear_q.state_dict(), state_path_4bit)

        size_orig, size_4 = (
            os.path.getsize(state_path),
            os.path.getsize(state_path_4bit),
        )
        size_ratio = size_4 / size_orig
        target_compression = (
            0.143 if original_dtype == torch.float32 else 0.29
        )  # these numbers get lower as weight shape increases
        ratio_error_msg = (
            f"quantized_size {size_4:,} is larger on disk than {target_compression:.2%} of original size {size_orig:,}"
        )
        assert size_ratio < target_compression, ratio_error_msg

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_copy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)

    shallow_copy_param = copy.copy(param)
    assert param.quant_state is shallow_copy_param.quant_state
    assert param.data.data_ptr() == shallow_copy_param.data.data_ptr()

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
def test_params4bit_torch_chunk_split(device, quant_type):
    """Test that torch.chunk and torch.split preserve Params4bit subclass for FSDP2 compatibility."""
    if device == "hpu" and not is_supported_on_hpu(quant_type, torch.float16, torch.uint8):
        pytest.skip("This configuration is not supported on HPU.")

    if device == "cpu":
        pytest.skip("CPU quantization causes segfault, skipping CPU test")

    original_tensor = torch.randn(8, 4, dtype=torch.float16, device="cpu")

    params4bit = bnb.nn.Params4bit(data=original_tensor, quant_type=quant_type, requires_grad=False)

    if device != "cpu":
        params4bit = params4bit.to(device)

    chunks = torch.chunk(params4bit, 2, dim=0)

    assert isinstance(chunks, tuple), "torch.chunk should return tuple"
    for chunk in chunks:
        assert isinstance(chunk, bnb.nn.Params4bit), "Chunk should preserve Params4bit subclass"
        assert hasattr(chunk, "quant_type"), "Should preserve metadata"
        assert chunk.quant_type == params4bit.quant_type, "Should preserve quant_type value"

    splits = torch.split(params4bit, 2, dim=0)

    assert isinstance(splits, tuple), "torch.split should return tuple"
    assert len(splits) > 0, "Should have at least one split"
    for split in splits:
        assert isinstance(split, bnb.nn.Params4bit), "Split should preserve Params4bit subclass"
        assert hasattr(split, "quant_type"), "Should preserve metadata"
        assert split.quant_type == params4bit.quant_type, "Should preserve quant_type value"

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_deepcopy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)
    dict_keys_before = set(param.__dict__.keys())
    copy_param = copy.deepcopy(param)
    dict_keys_after = set(param.__dict__.keys())
    dict_keys_copy = set(copy_param.__dict__.keys())

    assert param.quant_state is not copy_param.quant_state
    assert param.data.data_ptr() != copy_param.data.data_ptr()

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_copy

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_params4bit_real_serialization(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    original_tensor = torch.randn(300, 400)
    original_param = bnb.nn.Params4bit(
        data=original_tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
    )
    dict_keys_before = set(original_param.__dict__.keys())

    original_param.to(device)  # change device to trigger quantization

    serialized_param = pickle.dumps(original_param)
    deserialized_param = pickle.loads(serialized_param)
    dict_keys_after = set(original_param.__dict__.keys())
    dict_keys_deserialized = set(deserialized_param.__dict__.keys())

    assert torch.equal(original_param.data, deserialized_param.data)
    assert original_param.requires_grad == deserialized_param.requires_grad == False
    assert original_param.quant_type == deserialized_param.quant_type
    assert original_param.blocksize == deserialized_param.blocksize
    assert original_param.compress_statistics == deserialized_param.compress_statistics
    assert original_param.quant_state == deserialized_param.quant_state

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_deserialized

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("compute_dtype", [torch.bfloat16, torch.float32], ids=describe_dtype)
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
@pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph"))
@pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode"))
@pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4")
def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_statistics, bias, fullgraph, mode):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    if fullgraph and torch.__version__ < (2, 8, 0, "dev"):
        pytest.skip("fullgraph mode requires torch 2.8 or higher")

    if device == "cuda" and platform.system() == "Windows":
        pytest.skip("Triton is not officially supported on Windows")

    # Has a strange regression on Linux aarch64 CPU in torch==2.6.0 when fullgraph=False.
    if (
        not fullgraph
        and device == "cpu"
        and platform.machine() == "aarch64"
        and platform.system() == "Linux"
        and ((2, 7) > torch.__version__ >= (2, 6))
    ):
        pytest.xfail("Regression in torch==2.6.0 on Linux aarch64 CPU")

    dim = 256
    batch_size = 16

    torch.compiler.reset()

    # Create a small network with Linear4bit layers
    net = torch.nn.Sequential(
        *[
            bnb.nn.Linear4bit(
                dim,
                dim,
                bias=bias,
                compute_dtype=compute_dtype,
                compress_statistics=compress_statistics,
                quant_type=quant_type,
            )
            for _ in range(4)
        ]
    ).to(device)

    # Create input tensor
    x = torch.randn(batch_size, dim, dtype=compute_dtype, device=device)

    # Get reference output before compilation
    with torch.no_grad():
        ref_output = net(x)

    # Compile the model
    compile_backend = "hpu_backend" if device == "hpu" else "inductor"
    compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode, backend=compile_backend)

    # Get output from compiled model
    with torch.no_grad():
        compiled_output = compiled_net(x)

    # Check outputs match
    assert compiled_output.shape == ref_output.shape
    assert compiled_output.device == ref_output.device
    assert compiled_output.dtype == ref_output.dtype
    torch.testing.assert_close(compiled_output, ref_output)

    # Test with gradients
    x.requires_grad_(True)
    y1 = net(x).sum()
    y1.backward()
    grad_ref = x.grad.clone()

    x.grad = None
    y2 = compiled_net(x).sum()
    y2.backward()
    grad_compiled = x.grad.clone()

    torch.testing.assert_close(grad_compiled, grad_ref)


    is_supported_on_hpu,
    torch_load_from_buffer,
    torch_save_to_buffer,
)

storage = {
    "uint8": torch.uint8,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_storage", ["uint8", "float16", "bfloat16", "float32"])
@pytest.mark.parametrize("original_dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward"))
def test_linear_serialization(
    device, quant_type, original_dtype, compress_statistics, bias, quant_storage, save_before_forward
):
    if device == "hpu" and not is_supported_on_hpu(quant_type, original_dtype, storage[quant_storage]):
        pytest.skip("This configuration is not supported on HPU.")

    compute_dtype = None
    layer_shape = (300, 400)

    linear = torch.nn.Linear(*layer_shape, dtype=original_dtype, device="cpu")  # original layer

    # Quantizing original layer
    linear_q = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        device="meta",
    )
    new_weight = bnb.nn.Params4bit(data=linear.weight, quant_type=quant_type, requires_grad=False)
    linear_q.weight = new_weight
    if bias:
        linear_q.bias = torch.nn.Parameter(linear.bias)
    linear_q = linear_q.to(device)

    # saving to state_dict:
    sd = linear_q.state_dict()

    # restoring from state_dict:
    bias_data2 = sd.pop("bias", None)
    weight_data2 = sd.pop("weight")
    weight2 = bnb.nn.Params4bit.from_prequantized(quantized_stats=sd, data=weight_data2, device=device)

    # creating new layer with same params:
    linear_q2 = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        device="meta",
    )
    # loading weights from state_dict:
    linear_q2.weight = weight2
    if bias:
        linear_q2.bias = torch.nn.Parameter(bias_data2)
    linear_q2 = linear_q2.to(device)

    # MATCHING
    a, b = linear_q.weight, linear_q2.weight

    # Quantizing original layer with specified quant_storage type
    linear_qs = bnb.nn.Linear4bit(
        linear.in_features,
        linear.out_features,
        bias=bias,
        compute_dtype=compute_dtype,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
        quant_storage=storage[quant_storage],
        device="meta",
    )
    linear_qs.weight = bnb.nn.Params4bit(
        data=linear.weight,
        requires_grad=False,
        quant_type=quant_type,
        quant_storage=storage[quant_storage],
    )
    if bias:
        linear_qs.bias = torch.nn.Parameter(linear.bias)
    linear_qs = linear_qs.to(device)

    assert a.device == b.device
    assert a.dtype == b.dtype
    assert torch.equal(a, b)

    q0 = a.quant_state
    q1 = b.quant_state
    for attr in ("code", "dtype", "blocksize", "absmax"):
        c, d = getattr(q0, attr), getattr(q1, attr)
        if isinstance(c, torch.Tensor):
            assert torch.equal(c, d)
        else:
            assert c == d, f"{c} != {d}"

    if q0.state2 is not None:
        for attr in ("code", "dtype", "blocksize", "absmax"):
            c, d = getattr(q0.state2, attr), getattr(q1.state2, attr)
            if isinstance(c, torch.Tensor):
                assert torch.equal(c, d)
            else:
                assert c == d, f"{c} != {d}"

    if bias:
        a, b = linear_q.bias, linear_q2.bias
        assert a.device == b.device
        assert a.dtype == b.dtype
        assert torch.equal(a, b)

    if save_before_forward:
        bytes_4bit = torch_save_to_buffer(linear_q)

    # Forward test
    x = torch.rand(42, layer_shape[0], device=device)
    a = linear_q(x)
    b = linear_q2(x)
    c = linear_qs(x)
    assert a.device == b.device
    assert a.dtype == b.dtype
    assert a.device == c.device
    assert a.dtype == c.dtype
    assert torch.equal(a, b)
    assert torch.equal(a, c)

    if not save_before_forward:
        bytes_4bit = torch_save_to_buffer(linear_q)
    linear_q3 = torch_load_from_buffer(bytes_4bit)

    # Test moving to CPU and back to GPU
    if device != "cpu":
        linear_q2.to("cpu")
        linear_q2.to(device)
    d = linear_qs(x)
    assert c.dtype == d.dtype
    assert c.device == d.device
    assert torch.equal(c, d)

    d = linear_q3(x)
    assert c.dtype == d.dtype
    assert c.device == d.device
    assert torch.equal(c, d)

    # Saved size ratio test. Target set for layer_shape == (300, 400) w/ bias
    with TemporaryDirectory() as tmpdir:
        state_path_4bit = os.path.join(tmpdir, "state_4bit.pth")
        state_path = os.path.join(tmpdir, "state.pth")
        torch.save(linear.state_dict(), state_path)
        torch.save(linear_q.state_dict(), state_path_4bit)

        size_orig, size_4 = (
            os.path.getsize(state_path),
            os.path.getsize(state_path_4bit),
        )
        size_ratio = size_4 / size_orig
        target_compression = (
            0.143 if original_dtype == torch.float32 else 0.29
        )  # these numbers get lower as weight shape increases
        ratio_error_msg = (
            f"quantized_size {size_4:,} is larger on disk than {target_compression:.2%} of original size {size_orig:,}"
        )
        assert size_ratio < target_compression, ratio_error_msg

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_copy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)

    shallow_copy_param = copy.copy(param)


        torch.save(linear_q.state_dict(), state_path_4bit)

        size_orig, size_4 = (
            os.path.getsize(state_path),
            os.path.getsize(state_path_4bit),
        )
        size_ratio = size_4 / size_orig
        target_compression = (
            0.143 if original_dtype == torch.float32 else 0.29
        )  # these numbers get lower as weight shape increases
        ratio_error_msg = (
            f"quantized_size {size_4:,} is larger on disk than {target_compression:.2%} of original size {size_orig:,}"
        )
        assert size_ratio < target_compression, ratio_error_msg

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_copy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)

    shallow_copy_param = copy.copy(param)
    assert param.quant_state is shallow_copy_param.quant_state
    assert param.data.data_ptr() == shallow_copy_param.data.data_ptr()

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
def test_params4bit_torch_chunk_split(device, quant_type):
    """Test that torch.chunk and torch.split preserve Params4bit subclass for FSDP2 compatibility."""
    if device == "hpu" and not is_supported_on_hpu(quant_type, torch.float16, torch.uint8):
        pytest.skip("This configuration is not supported on HPU.")

    if device == "cpu":
        pytest.skip("CPU quantization causes segfault, skipping CPU test")

    original_tensor = torch.randn(8, 4, dtype=torch.float16, device="cpu")

    params4bit = bnb.nn.Params4bit(data=original_tensor, quant_type=quant_type, requires_grad=False)

    if device != "cpu":
        params4bit = params4bit.to(device)

    chunks = torch.chunk(params4bit, 2, dim=0)


def test_copy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)

    shallow_copy_param = copy.copy(param)
    assert param.quant_state is shallow_copy_param.quant_state
    assert param.data.data_ptr() == shallow_copy_param.data.data_ptr()

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
def test_params4bit_torch_chunk_split(device, quant_type):
    """Test that torch.chunk and torch.split preserve Params4bit subclass for FSDP2 compatibility."""
    if device == "hpu" and not is_supported_on_hpu(quant_type, torch.float16, torch.uint8):
        pytest.skip("This configuration is not supported on HPU.")

    if device == "cpu":
        pytest.skip("CPU quantization causes segfault, skipping CPU test")

    original_tensor = torch.randn(8, 4, dtype=torch.float16, device="cpu")

    params4bit = bnb.nn.Params4bit(data=original_tensor, quant_type=quant_type, requires_grad=False)

    if device != "cpu":
        params4bit = params4bit.to(device)

    chunks = torch.chunk(params4bit, 2, dim=0)

    assert isinstance(chunks, tuple), "torch.chunk should return tuple"
    for chunk in chunks:
        assert isinstance(chunk, bnb.nn.Params4bit), "Chunk should preserve Params4bit subclass"
        assert hasattr(chunk, "quant_type"), "Should preserve metadata"
        assert chunk.quant_type == params4bit.quant_type, "Should preserve quant_type value"

    splits = torch.split(params4bit, 2, dim=0)

    assert isinstance(splits, tuple), "torch.split should return tuple"
    assert len(splits) > 0, "Should have at least one split"
    for split in splits:
        assert isinstance(split, bnb.nn.Params4bit), "Split should preserve Params4bit subclass"
        assert hasattr(split, "quant_type"), "Should preserve metadata"
        assert split.quant_type == params4bit.quant_type, "Should preserve quant_type value"

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_deepcopy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)
    dict_keys_before = set(param.__dict__.keys())
    copy_param = copy.deepcopy(param)


    assert isinstance(chunks, tuple), "torch.chunk should return tuple"
    for chunk in chunks:
        assert isinstance(chunk, bnb.nn.Params4bit), "Chunk should preserve Params4bit subclass"
        assert hasattr(chunk, "quant_type"), "Should preserve metadata"
        assert chunk.quant_type == params4bit.quant_type, "Should preserve quant_type value"

    splits = torch.split(params4bit, 2, dim=0)

    assert isinstance(splits, tuple), "torch.split should return tuple"
    assert len(splits) > 0, "Should have at least one split"
    for split in splits:
        assert isinstance(split, bnb.nn.Params4bit), "Split should preserve Params4bit subclass"
        assert hasattr(split, "quant_type"), "Should preserve metadata"
        assert split.quant_type == params4bit.quant_type, "Should preserve quant_type value"

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_deepcopy_param(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    tensor = torch.randn(300, 400)
    param = bnb.nn.Params4bit(
        data=tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)
    dict_keys_before = set(param.__dict__.keys())
    copy_param = copy.deepcopy(param)
    dict_keys_after = set(param.__dict__.keys())
    dict_keys_copy = set(copy_param.__dict__.keys())

    assert param.quant_state is not copy_param.quant_state
    assert param.data.data_ptr() != copy_param.data.data_ptr()

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_copy

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_params4bit_real_serialization(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    original_tensor = torch.randn(300, 400)
    original_param = bnb.nn.Params4bit(
        data=original_tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
    )
    dict_keys_before = set(original_param.__dict__.keys())

    original_param.to(device)  # change device to trigger quantization


        compress_statistics=compress_statistics,
        requires_grad=False,
    ).to(device)
    dict_keys_before = set(param.__dict__.keys())
    copy_param = copy.deepcopy(param)
    dict_keys_after = set(param.__dict__.keys())
    dict_keys_copy = set(copy_param.__dict__.keys())

    assert param.quant_state is not copy_param.quant_state
    assert param.data.data_ptr() != copy_param.data.data_ptr()

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_copy

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128])
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
def test_params4bit_real_serialization(device, quant_type, blocksize, compress_statistics):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    original_tensor = torch.randn(300, 400)
    original_param = bnb.nn.Params4bit(
        data=original_tensor,
        quant_type=quant_type,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
    )
    dict_keys_before = set(original_param.__dict__.keys())

    original_param.to(device)  # change device to trigger quantization

    serialized_param = pickle.dumps(original_param)
    deserialized_param = pickle.loads(serialized_param)
    dict_keys_after = set(original_param.__dict__.keys())
    dict_keys_deserialized = set(deserialized_param.__dict__.keys())

    assert torch.equal(original_param.data, deserialized_param.data)
    assert original_param.requires_grad == deserialized_param.requires_grad == False
    assert original_param.quant_type == deserialized_param.quant_type
    assert original_param.blocksize == deserialized_param.blocksize
    assert original_param.compress_statistics == deserialized_param.compress_statistics
    assert original_param.quant_state == deserialized_param.quant_state

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_deserialized

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("compute_dtype", [torch.bfloat16, torch.float32], ids=describe_dtype)
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
@pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph"))
@pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode"))
@pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4")
def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_statistics, bias, fullgraph, mode):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    if fullgraph and torch.__version__ < (2, 8, 0, "dev"):
        pytest.skip("fullgraph mode requires torch 2.8 or higher")

    if device == "cuda" and platform.system() == "Windows":
        pytest.skip("Triton is not officially supported on Windows")


    assert torch.equal(original_param.data, deserialized_param.data)
    assert original_param.requires_grad == deserialized_param.requires_grad == False
    assert original_param.quant_type == deserialized_param.quant_type
    assert original_param.blocksize == deserialized_param.blocksize
    assert original_param.compress_statistics == deserialized_param.compress_statistics
    assert original_param.quant_state == deserialized_param.quant_state

    # there was a bug where deepcopy would modify the original object
    assert dict_keys_before == dict_keys_after
    assert dict_keys_before == dict_keys_deserialized

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
@pytest.mark.parametrize("compute_dtype", [torch.bfloat16, torch.float32], ids=describe_dtype)
@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
@pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph"))
@pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode"))
@pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4")
def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_statistics, bias, fullgraph, mode):
    if device == "hpu" and not is_supported_on_hpu(quant_type):
        pytest.skip("This configuration is not supported on HPU.")

    if fullgraph and torch.__version__ < (2, 8, 0, "dev"):
        pytest.skip("fullgraph mode requires torch 2.8 or higher")

    if device == "cuda" and platform.system() == "Windows":
        pytest.skip("Triton is not officially supported on Windows")

    # Has a strange regression on Linux aarch64 CPU in torch==2.6.0 when fullgraph=False.
    if (
        not fullgraph
        and device == "cpu"
        and platform.machine() == "aarch64"
        and platform.system() == "Linux"
        and ((2, 7) > torch.__version__ >= (2, 6))
    ):
        pytest.xfail("Regression in torch==2.6.0 on Linux aarch64 CPU")

    dim = 256
    batch_size = 16

    torch.compiler.reset()

    # Create a small network with Linear4bit layers
    net = torch.nn.Sequential(
        *[
            bnb.nn.Linear4bit(
                dim,
                dim,
                bias=bias,
                compute_dtype=compute_dtype,
                compress_statistics=compress_statistics,
                quant_type=quant_type,
            )
            for _ in range(4)
        ]
    ).to(device)

    # Create input tensor
    x = torch.randn(batch_size, dim, dtype=compute_dtype, device=device)

    # Get reference output before compilation
    with torch.no_grad():
        ref_output = net(x)

    # Compile the model
    compile_backend = "hpu_backend" if device == "hpu" else "inductor"
    compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode, backend=compile_backend)

    # Get output from compiled model
    with torch.no_grad():
        compiled_output = compiled_net(x)

    # Check outputs match
    assert compiled_output.shape == ref_output.shape
    assert compiled_output.device == ref_output.device
    assert compiled_output.dtype == ref_output.dtype
    torch.testing.assert_close(compiled_output, ref_output)

    # Test with gradients
    x.requires_grad_(True)
    y1 = net(x).sum()
    y1.backward()
    grad_ref = x.grad.clone()

    x.grad = None
    y2 = compiled_net(x).sum()
    y2.backward()
    grad_compiled = x.grad.clone()

    torch.testing.assert_close(grad_compiled, grad_ref)


import torch
import torch._dynamo
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# torch._dynamo.config.suppress_errors = True

torch.set_float32_matmul_precision("high")

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

# torch._dynamo.config.capture_dynamic_output_shape_ops = True

model_id = "google/gemma-2-2b-it"
# model_id = "Qwen/Qwen2.5-7B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

input_text = "Write me a poem about Machine Learning."
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)

# model.forward = torch.compile(model.forward, fullgraph=True)

model = torch.compile(model)

outputs = model.generate(**input_ids, max_new_tokens=32)
print(tokenizer.decode(outputs[0]))


import torch
from transformers import LlamaForCausalLM, LlamaTokenizer

MAX_NEW_TOKENS = 128
model_name = "meta-llama/Llama-2-7b-hf"

text = "Hamburg is in which country?\n"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
input_ids = tokenizer(text, return_tensors="pt").input_ids

max_memory = f"{int(torch.cuda.mem_get_info()[0] / 1024**3) - 2}GB"

n_gpus = torch.cuda.device_count()
max_memory = {i: max_memory for i in range(n_gpus)}

model = LlamaForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True, max_memory=max_memory)

generated_ids = model.generate(input_ids, max_length=MAX_NEW_TOKENS)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))


import dataclasses
from functools import lru_cache
import logging
import re
import subprocess
from typing import Optional

import torch

@dataclasses.dataclass(frozen=True)
class CUDASpecs:
    highest_compute_capability: tuple[int, int]
    cuda_version_string: str
    cuda_version_tuple: tuple[int, int]

    @property
    def has_imma(self) -> bool:
        return torch.version.hip or self.highest_compute_capability >= (7, 5)

def get_compute_capabilities() -> list[tuple[int, int]]:
    return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count()))

@lru_cache(None)
def get_cuda_version_tuple() -> Optional[tuple[int, int]]:
    """Get CUDA/HIP version as a tuple of (major, minor)."""
    try:
        if torch.version.cuda:
            version_str = torch.version.cuda
        elif torch.version.hip:
            version_str = torch.version.hip
        else:
            return None

        parts = version_str.split(".")
        if len(parts) >= 2:
            return tuple(map(int, parts[:2]))
        return None
    except (AttributeError, ValueError, IndexError):
        return None

def get_cuda_version_string() -> Optional[str]:
    """Get CUDA/HIP version as a string."""
    version_tuple = get_cuda_version_tuple()
    if version_tuple is None:
        return None
    major, minor = version_tuple
    return f"{major * 10 + minor}"

def get_cuda_specs() -> Optional[CUDASpecs]:
    """Get CUDA/HIP specifications."""
    if not torch.cuda.is_available():
        return None

    try:
        compute_capabilities = get_compute_capabilities()
        if not compute_capabilities:
            return None

        version_tuple = get_cuda_version_tuple()
        if version_tuple is None:
            return None

        version_string = get_cuda_version_string()
        if version_string is None:
            return None

        return CUDASpecs(
            highest_compute_capability=compute_capabilities[-1],
            cuda_version_string=version_string,
            cuda_version_tuple=version_tuple,
        )
    except Exception:
        return None

def get_rocm_gpu_arch() -> str:
    """Get ROCm GPU architecture."""
    logger = logging.getLogger(__name__)
    try:
        if torch.version.hip:
            result = subprocess.run(["rocminfo"], capture_output=True, text=True)
            match = re.search(r"Name:\s+gfx([a-zA-Z\d]+)", result.stdout)
            if match:
                return "gfx" + match.group(1)
            else:
                return "unknown"
        else:
            return "unknown"
    except Exception as e:
        logger.error(f"Could not detect ROCm GPU architecture: {e}")
        if torch.cuda.is_available():
            logger.warning(
                """
ROCm GPU architecture detection failed despite ROCm being available.
                """,
            )
        return "unknown"


CUDASpecs

import dataclasses
from functools import lru_cache
import logging
import re
import subprocess
from typing import Optional

import torch

@dataclasses.dataclass(frozen=True)
class CUDASpecs:
    highest_compute_capability: tuple[int, int]
    cuda_version_string: str
    cuda_version_tuple: tuple[int, int]

    @property
    def has_imma(self) -> bool:
        return torch.version.hip or self.highest_compute_capability >= (7, 5)

def get_compute_capabilities() -> list[tuple[int, int]]:
    return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count()))

@lru_cache(None)
def get_cuda_version_tuple() -> Optional[tuple[int, int]]:
    """Get CUDA/HIP version as a tuple of (major, minor)."""
    try:
        if torch.version.cuda:
            version_str = torch.version.cuda
        elif torch.version.hip:
            version_str = torch.version.hip
        else:
            return None

        parts = version_str.split(".")
        if len(parts) >= 2:
            return tuple(map(int, parts[:2]))


from functools import lru_cache
import logging
import re
import subprocess
from typing import Optional

import torch

@dataclasses.dataclass(frozen=True)
class CUDASpecs:
    highest_compute_capability: tuple[int, int]
    cuda_version_string: str
    cuda_version_tuple: tuple[int, int]

    @property
    def has_imma(self) -> bool:
        return torch.version.hip or self.highest_compute_capability >= (7, 5)

def get_compute_capabilities() -> list[tuple[int, int]]:
    return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count()))

@lru_cache(None)
def get_cuda_version_tuple() -> Optional[tuple[int, int]]:
    """Get CUDA/HIP version as a tuple of (major, minor)."""
    try:
        if torch.version.cuda:
            version_str = torch.version.cuda
        elif torch.version.hip:
            version_str = torch.version.hip
        else:
            return None

        parts = version_str.split(".")
        if len(parts) >= 2:
            return tuple(map(int, parts[:2]))
        return None
    except (AttributeError, ValueError, IndexError):
        return None


import torch

@dataclasses.dataclass(frozen=True)
class CUDASpecs:
    highest_compute_capability: tuple[int, int]
    cuda_version_string: str
    cuda_version_tuple: tuple[int, int]

    @property
    def has_imma(self) -> bool:
        return torch.version.hip or self.highest_compute_capability >= (7, 5)

def get_compute_capabilities() -> list[tuple[int, int]]:
    return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count()))

@lru_cache(None)
def get_cuda_version_tuple() -> Optional[tuple[int, int]]:
    """Get CUDA/HIP version as a tuple of (major, minor)."""
    try:
        if torch.version.cuda:
            version_str = torch.version.cuda
        elif torch.version.hip:
            version_str = torch.version.hip
        else:
            return None

        parts = version_str.split(".")
        if len(parts) >= 2:
            return tuple(map(int, parts[:2]))
        return None
    except (AttributeError, ValueError, IndexError):
        return None

def get_cuda_version_string() -> Optional[str]:
    """Get CUDA/HIP version as a string."""
    version_tuple = get_cuda_version_tuple()
    if version_tuple is None:
        return None
    major, minor = version_tuple
    return f"{major * 10 + minor}"

def get_cuda_specs() -> Optional[CUDASpecs]:
    """Get CUDA/HIP specifications."""
    if not torch.cuda.is_available():
        return None

    try:
        compute_capabilities = get_compute_capabilities()
        if not compute_capabilities:
            return None