| 124 | |
| 125 | |
| 126 | def test(): |
| 127 | import torch |
| 128 | torch.manual_seed(1234) |
| 129 | for rep in range(4): |
| 130 | length: int = torch.randint(10, 2_000, (1, )).item() |
| 131 | bits: int = torch.randint(1, 16, (1, )).item() |
| 132 | tokens: tp.List[int] = torch.randint(2**bits, (length, )).tolist() |
| 133 | rebuilt: tp.List[int] = [] |
| 134 | buf = io.BytesIO() |
| 135 | packer = BitPacker(bits, buf) |
| 136 | for token in tokens: |
| 137 | packer.push(token) |
| 138 | packer.flush() |
| 139 | buf.seek(0) |
| 140 | unpacker = BitUnpacker(bits, buf) |
| 141 | while True: |
| 142 | value = unpacker.pull() |
| 143 | if value is None: |
| 144 | break |
| 145 | rebuilt.append(value) |
| 146 | assert len(rebuilt) >= len(tokens), (len(rebuilt), len(tokens)) |
| 147 | # The flushing mechanism might lead to "ghost" values at the end of the stream. |
| 148 | assert len(rebuilt) <= len(tokens) + 8 // bits, (len(rebuilt), |
| 149 | len(tokens), bits) |
| 150 | for idx, (a, b) in enumerate(zip(tokens, rebuilt)): |
| 151 | assert a == b, (idx, a, b) |
| 152 | |
| 153 | |
| 154 | if __name__ == '__main__': |