Class TestNonContiguousInputs

tests/test_ops.py:378–474

Regression tests for #1342 and #1690: quantization must handle non-contiguous tensors correctly.

class TestNonContiguousInputs:
    """Regression tests for #1342 and #1690: quantization must handle non-contiguous tensors correctly."""

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
    @pytest.mark.parametrize("blocksize", [64, 128, 256])
    def test_quantize_blockwise_non_contiguous(self, device, dtype, blocksize):
        code = bitsandbytes.functional.create_dynamic_map().to(device)

        # Create non-contiguous tensor via slicing
        A_full = torch.randn(3, 4, 6, 256, dtype=dtype, device=device)
        A_noncontig = A_full[:, ::2, :, :]
        assert not A_noncontig.is_contiguous()

        A_contig = A_noncontig.contiguous()

        out_nc, absmax_nc = torch.ops.bitsandbytes.quantize_blockwise(A_noncontig, code, blocksize)
        out_c, absmax_c = torch.ops.bitsandbytes.quantize_blockwise(A_contig, code, blocksize)

        torch.testing.assert_close(absmax_nc, absmax_c)
        torch.testing.assert_close(out_nc, out_c)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
    @pytest.mark.parametrize("blocksize", [64, 128, 256])
    def test_dequantize_blockwise_non_contiguous(self, device, dtype, blocksize):
        code = bitsandbytes.functional.create_dynamic_map().to(device, dtype=torch.float32)

        # Quantize a contiguous tensor, then create non-contiguous uint8 via transpose
        A = torch.randn(1024, 1024, dtype=dtype, device=device)
        quantized, absmax = torch.ops.bitsandbytes.quantize_blockwise(A, code, blocksize)

        # Create non-contiguous uint8 tensor by transposing and transposing back
        q_noncontig = quantized.t().t()
        # If that's still contiguous, use a different approach
        if q_noncontig.is_contiguous():
            # Pad and slice to force non-contiguity
            q_padded = torch.zeros(1024, 1025, dtype=torch.uint8, device=device)
            q_padded[:, :1024] = quantized
            q_noncontig = q_padded[:, :1024]

        assert not q_noncontig.is_contiguous()
        q_contig = q_noncontig.contiguous()

        out_nc = torch.ops.bitsandbytes.dequantize_blockwise(q_noncontig, absmax, code, blocksize, dtype)
        out_c = torch.ops.bitsandbytes.dequantize_blockwise(q_contig, absmax, code, blocksize, dtype)

        torch.testing.assert_close(out_nc, out_c)

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
    @pytest.mark.parametrize("quant_type", ["fp4", "nf4"])
    @pytest.mark.parametrize("blocksize", [64, 128, 256])
    def test_quantize_4bit_non_contiguous(self, device, dtype, quant_type, blocksize):
        if device not in ("cuda", "mps"):
            pytest.skip("Non-contiguous input handling not implemented for this backend")

        # Reproduce issue #1342: non-contiguous tensor from slicing