Skip to content

Commit 5fb32e0

Browse files
committed
Add example to README
1 parent 9d83231 commit 5fb32e0

File tree

1 file changed

+27
-1
lines changed

1 file changed

+27
-1
lines changed

README.md

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,33 @@ To install pi-quant from PyPI, run the following command:
3737
pip 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

Comments
 (0)