[dexhunter]

TL;DR

• This PR makes EntropyBottleneck.decompress faster, mainly the loop around decode_with_indexes in EntropyModel.decompress.
• It converts the CDF tables to lists once and uses numpy to build the decoded tensor.
• Measured on a workload like test_compression_ND, median of 3 runs: 68.6 ms -> 47.4 ms, 31% faster.
• The model, helper, entropy model and checkerboard tests pass: 113 passed, 3 skipped (same as master).

What & Why

Before, EntropyModel.decompress called cdf.tolist(), cdf_length.tolist() and offset.tolist() again for each string. These tables are the same for all strings. Then it built each output row with torch.tensor(values), and torch.tensor is slow on a long Python list of ints.

Now the code converts the three tables once, before the loop. It reads the decoded list with np.fromiter(..., dtype=np.int32) and uses torch.from_numpy to wrap it. For outputs with 2 or more spatial dims, it builds the indexes over one flat spatial dim. It also runs dequantize on a flat (N, C, -1) view and reshapes the result back. The values, dtype and shape stay the same. Only the number of Python and torch calls goes down.

Changes

• EntropyModel.decompress: convert cdf, cdf_length and offset to lists once, outside the loop.
• EntropyModel.decompress: use np.fromiter + torch.from_numpy instead of torch.tensor on the decoded list.
• EntropyModel.decompress: for inputs with 2 or more spatial dims, dequantize on a flat (N, C, -1) view.
• EntropyBottleneck.decompress: build the indexes over one flat spatial dim, then reshape to the full size.

Results

EntropyBottleneck.decompress time per call, median of 7 runs, one CPU thread:

Situation	Before	After	Change
batch 1, 128 ch, 0D (no spatial dims)	0.186 ms	0.183 ms	no change
batch 1, 128 ch, 16x16	2.78 ms	1.91 ms	31% faster
batch 1, 192 ch, 64x64	63.57 ms	41.23 ms	35% faster
batch 4, 128 ch, 32x32	42.65 ms	28.02 ms	34% faster
batch 2, 128 ch, 4x4x4x4x4	21.09 ms	13.98 ms	34% faster

Why this is safe and easy to maintain

The change is small (about 40 lines in one file) and keeps the same loop structure. It does not change public signatures or the C++ coders. The decoded values, dtype and shape stay the same, and the existing round-trip tests in tests/test_entropy_models.py and tests/test_models.py cover this path.

Validation

• Tests: python -m pytest tests/test_models.py tests/test_models_helpers.py tests/test_entropy_models.py tests/test_multi_context_checkerboard.py gives 113 passed, 3 skipped on both master and this branch. The skips are tests that need pytorch_wavelets.
• Benchmark: a workload like TestEntropyBottleneck::test_compression_ND (EntropyBottleneck(128), inputs with 0 to 5 spatial dims) takes 68.6 ms before and 47.4 ms after (3 runs: 0.690, 0.697, 0.690 of the old time).
• Exact output: the decoded tensors were compared byte for byte with master on the same inputs.
• ruff format --check and ruff check pass on the changed file.
• Short reproduction:

import time, torch
from compressai.entropy_models import EntropyBottleneck

eb = EntropyBottleneck(192)
eb.update(force=True)
x = torch.randn(1, 192, 64, 64) * 3
s = eb.compress(x)
t = time.perf_counter()
for _ in range(20):
    eb.decompress(s, x.shape[2:])
print((time.perf_counter() - t) / 20 * 1e3, "ms")

AI assistance

This change was found with AI help: Codex (gpt-6.1-sol) suggested candidates and Weco Observe tracked them. Automated checks picked one, then measured it and ran the tests listed above. Before you merge, please review the diff as you would any other contribution.

Weco share link: https://dashboard.weco.ai/share/Yphl6lyOc4sqZEzFk2UP_Pw_QAb4A_xH