@@ -37,7 +37,33 @@ To install pi-quant from PyPI, run the following command:
3737pip install pypiquant
3838```
3939
40- # Benchmarks
40+ ## Examples
41+ piquant is torch compatible. Here are some examples of how to use it with PyTorch:
42+
43+ ``` python
44+ import torch
45+ import piquant
46+
47+ # Quantize and back with: bfloat16 -> uint4 -> bfloat16
48+ # In torch, quint4x2 means two 4-bit quantized integers per byte.
49+ tensor = torch.rand(1000 , dtype = torch.bfloat16, device = ' cpu' )
50+
51+ # Compute quantization parameters for uint4 (needed for quantization and dequantization)
52+ scale, zero_point = piquant.torch.compute_quant_params(tensor, dtype = torch.quint4x2)
53+
54+ # Quantize the tensor to uint4
55+ quantized = piquant.torch.quantize(tensor, scale = scale, zero_point = zero_point, dtype = torch.quint4x2)
56+
57+ # Dequantize back to bfloat16
58+ dequantized = piquant.torch.dequantize(quantized, scale = scale, zero_point = zero_point, dtype = torch.bfloat16)
59+
60+ # Check if the dequantized tensor is close to the original tensor
61+ assert torch.allclose(dequantized, tensor, atol = scale* 0.5 + 1e-3 ), " Dequantization did not match original tensor"
62+
63+ # Print parts of original and dequantized tensors for verification
64+ print (" Original tensor (first 10 elements):" , tensor[:10 ].tolist())
65+ print (" Dequant tensor (first 10 elements):" , dequantized[:10 ].tolist())
66+ ```
4167
4268## Benchmark
4369
0 commit comments