| 91 | |
| 92 | @dataclass(frozen=True) |
| 93 | class Q8_0QuantizedDataType(QuantizedDataType): |
| 94 | # Mini Q8_0 quantization in Python! |
| 95 | def quantize(self, arr: NDArray) -> NDArray: |
| 96 | assert arr.size % self.block_size == 0 and arr.size != 0, f'Bad array size {arr.size}' |
| 97 | assert arr.dtype == np.float32, f'Bad array type {arr.dtype}' |
| 98 | n_blocks = arr.size // self.block_size |
| 99 | blocks = arr.reshape((n_blocks, self.block_size)) |
| 100 | # Much faster implementation of block quantization contributed by @Cebtenzzre |
| 101 | |
| 102 | def quantize_blocks_q8_0(blocks: NDArray) -> Iterable[tuple[Any, Any]]: |
| 103 | d = abs(blocks).max(axis = 1) / np.float32(127) |
| 104 | with np.errstate(divide = 'ignore'): |
| 105 | qs = (blocks / d[:, None]).round() |
| 106 | qs[d == 0] = 0 |
| 107 | yield from zip(d, qs) |
| 108 | return np.fromiter(quantize_blocks_q8_0(blocks), count = n_blocks, dtype = self.quantized_dtype) |
| 109 | |
| 110 | # @dataclass(frozen=True) |
| 111 | # class TransformedDataType(DataType): |