repack.py

import numpy as np


def repack(array, n):
    """Repack a uint8 array as a contiguous stream of n-bit values
    
    Works for arbitrary 0 < n <= 64, drops tail elements as necessary to evenly divide
    """
    if not 0 < n <= 64:
        raise ValueError(n)
    assert array.dtype == np.uint8
    if n in (8, 16, 32, 64):
        x = array
        tail = len(array) % (n//8)
        if tail != 0:
            x = x[:-tail]
        return x.view(f">u{n//8}")
    x = np.unpackbits(array)
    tail = (len(array) * 8) % n
    if tail != 0:
        x = x[:-tail]
    x = np.packbits(x.reshape(-1, n), axis=1)
    minwidth = 2 ** max(0, int(np.ceil(np.log2(n / 8))))
    x = np.pad(x, ((0, 0), (minwidth - x.shape[-1], 0))).view(f">u{minwidth}")
    assert x.shape[-1] == 1
    return x[:, 0] >> (8 * minwidth - n)