flashinfer.quantization.segment_packbits#

flashinfer.quantization.segment_packbits(x: torch.Tensor, indptr: torch.Tensor, bitorder: str = 'big')#

Pack a batch of binary-valued segments into bits in a uint8 array.

For each segment, the semantics of this function is the same as numpy.packbits.

Parameters:

x (torch.Tensor) – The 1D binary-valued array to pack, shape (indptr[-1],).
indptr (torch.Tensor) – The index pointer of each segment in x, shape (batch_size + 1,). The i-th segment in x is x[indptr[i]:indptr[i+1]].
bitorder (str) – The bit-order (“bit”/”little”) of the output. Default is “big”.

Returns:

y (torch.Tensor) – An uint8 packed array, shape: (new_indptr[-1],). The y[new_indptr[i]:new_indptr[i+1]] contains the packed bits x[indptr[i]:indptr[i+1]].
new_indptr (torch.Tensor) – The new index pointer of each packed segment in y, shape (batch_size + 1,). It’s guaranteed that new_indptr[i+1] - new_indptr[i] == (indptr[i+1] - indptr[i] + 7) // 8.