feat(entropy): entropy-coding research framework for RVQ/SPIHT priors#62
Merged
Conversation
…ithmetic benchmarks) Checkpoint of in-progress work handed off between agent sessions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Fixes for PR readiness: - compressionkit/dsp/spiht.py: add missing `Any` import (F821 — used in the new sink/source type hints, previously undefined at lint time). - Apply ruff --fix for unused imports/noqa, RUF005/RUF046 style fixes. - Apply ruff format to bring 6 files (written without the formatter) into compliance with the repo formatting convention. No behavioral changes; full test suite (488 tests) still passes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Adds a research-oriented, pluggable entropy-coding layer that can sit losslessly on top of existing frozen codecs (RVQ/SPIHT/Hybrid), including learned priors, two backend bit-packers (arithmetic + rANS), SPIHT sink/source injection for neural bit priors, and a benchmark harness + documentation to support repeatable experiments without touching golden/production defaults.
Changes:
- Introduces
EntropyCoder-conforming implementations (learned + classical baselines) plus an evaluation harness for scorecard-style comparisons and noise sweeps. - Extends the two-stage runtime with correctness fixes (CDF quantization) and a new
"rans"backend, and updates prior APIs to align with sliding-context incremental inference. - Adds SPIHT neural bit-prior numpy inference + injectable sink/source seam, along with extensive tests, research scripts, and documentation.
Reviewed changes
Copilot reviewed 20 out of 20 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| zensical.toml | Adds the entropy-coding research doc to the documentation nav. |
| docs/entropy-coding-research.md | Documents the entropy-coding research “blocks → experiments → golden promotion” workflow and known pitfalls. |
| AGENTS.md | Adds explicit guidance to reuse existing compressionkit/ blocks and avoid duplicated script logic. |
| compressionkit/runtime/entropy_algorithms.py | Adds pluggable entropy coders (learned prior wrapper + static/Markov baselines) using arithmetic/rANS backends. |
| compressionkit/evaluation/entropy_benchmark.py | Adds a benchmark harness to evaluate entropy coders on token streams and across SNR/noise conditions. |
| compressionkit/runtime/prior.py | Adds predict_next_probs() and updates predict_log_probs() to use sliding context/incremental convention. |
| compressionkit/runtime/two_stage.py | Adds "rans" backend, fixes CDF quantization total to avoid zero-width intervals, and aligns probability extraction with incremental context. |
| compressionkit/dsp/spiht.py | Adds generic sink=/source= dependency injection seam and bitplane hook support for neural entropy coding. |
| compressionkit/generative/causal_priors.py | Adds causal prior architectures and SplitHalf layer to avoid Lambda-serialization fragility; includes SPIHT bit-prior builder. |
| compressionkit/generative/spiht_bit_prior.py | Adds pure-numpy WaveNet bit-prior inference plus NeuralAcSink/NeuralAcSource and incremental predictor. |
| compressionkit/generative/init.py | Re-exports causal-prior builders and SplitHalf from the package namespace. |
| scripts/measure_rvq_entropy.py | Adds LR schedule + steps_per_epoch support and dataset-repeat training flow for prior training experiments. |
| scripts/measure_rvq_entropy_2level.py | Adds a 2-level (time × level) RVQ prior experiment that avoids token interleaving. |
| scripts/benchmark_rans_vs_arithmetic.py | Adds an experiment to compare actual rANS vs arithmetic bit usage using trained priors and cached token streams. |
| scripts/build_augmented_ecg_train_cache.py | Adds a script to build an augmented ECG token cache by reusing the codec’s own training pipeline/augmentation. |
| tests/test_two_stage.py | Adds regression coverage for the CDF zero-width bug and new rANS backend roundtrips + backend selection errors. |
| tests/test_entropy_algorithms.py | Adds protocol/roundtrip tests for learned/static/Markov entropy coders across arithmetic/rANS. |
| tests/test_entropy_benchmark.py | Adds tests for the benchmark harness and noise sweep behavior. |
| tests/test_causal_priors.py | Adds save/load and trainability regression tests for WaveNet priors (avoids Lambda deserialization breakage). |
| tests/test_spiht_bit_prior.py | Adds correctness + integration tests for SPIHT neural bit prior (numpy forward pass, incremental predictor, sink/source roundtrip). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
175
to
179
| @@ -149,11 +178,15 @@ | |||
| indices_4d = indices[np.newaxis] | |||
| bpt_prior = self._prior.bits_per_token(indices_4d) | |||
Comment on lines
+636
to
+642
| slot = x & (_RANS_TOTAL - 1) | ||
| sym = 0 | ||
| for s in range(vocab_size): | ||
| if starts[s] <= slot < starts[s] + freqs[s]: | ||
| sym = s | ||
| break | ||
| decoded_tokens[t] = sym |
Comment on lines
+528
to
533
| if sink is not None: | ||
| if log_emissions: | ||
| raise ValueError("sink and log_emissions are mutually exclusive.") | ||
| elif log_emissions: | ||
| if not use_ac: | ||
| raise ValueError("log_emissions=True requires use_ac=True") |
Comment on lines
+145
to
+147
| def predict_next_probs(self, context_tokens: np.ndarray) -> np.ndarray: | ||
| batch = context_tokens.shape[0] | ||
| return np.tile(self._probs, (batch, 1)) |
Comment on lines
+187
to
+195
| # Level-1 head sees h_t (causal state) AND the *next* step's true level0 | ||
| # embedding (teacher forcing) — same trick a coarse-to-fine / residual | ||
| # token model uses. Pad the shift with a learned "unknown" embedding so | ||
| # shapes stay static (last position has no "next" ground truth). | ||
| next_level0_embed = embed0(level0_in) # reuse table; shift below | ||
| shifted = keras.layers.Lambda( | ||
| lambda t: keras.ops.pad(t[:, 1:, :], [[0, 0], [0, 1], [0, 0]]), | ||
| name="shift_level0_embed", | ||
| )(next_level0_embed) |
apage224
marked this pull request as ready for review
July 5, 2026 20:15
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f0c0287565
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| x1 = np.stack([level1[s : s + context_length] for s in bs]).astype(np.int32) | ||
| logits0, logits1 = model([x0, x1], training=False) | ||
| last0 = keras.ops.convert_to_numpy(logits0[:, -1, :]).astype(np.float64) | ||
| last1 = keras.ops.convert_to_numpy(logits1[:, -1, :]).astype(np.float64) |
There was a problem hiding this comment.
When running the 2-level entropy script, level1_logits[:, -1, :] scores the only position where shift_level0_embed has padded away the “next” level-0 embedding, so every validation level-1 token is measured without the same-step level-0 code that the model is trained and documented to use. This makes entropy_report.json report a different, weaker distribution for level-1 and can invalidate comparisons against the 1-D prior; score a position whose shifted level-0 token is present, or append/provide the target level-0 before evaluating level-1.
Useful? React with 👍 / 👎.
Addresses Copilot/Codex review findings on PR #62: - two_stage.py: compress_indices() no longer does a second full prior pass to compute bits_per_token — derives it directly from the already-computed probs_per_position (same windowed-context values), halving TFLite interpreter calls per encode. - two_stage.py: rANS decode now finds each token via binary search (bisect) over the strictly-increasing cumulative starts array instead of an O(vocab_size) linear scan per token. - spiht.py: injecting a custom sink now requires use_ac=True (previously only documented, not enforced) — prevents silently producing an undecodable bitstream whose metadata says use_ac=False. - entropy_algorithms.py: fixed _StaticPriorAdapter docstring to match actual behavior (it always returns the fitted histogram, including at position 0 — the callers, not this adapter, are what special-case position 0 elsewhere). - measure_rvq_entropy_2level.py: replaced a closure-based keras.layers.Lambda with Cropping1D + ZeroPadding1D (the same fragile-serialization pattern this PR introduced SplitHalf to avoid elsewhere). - measure_rvq_entropy_2level.py: fixed a real evaluation bug — the entropy eval scored level-1 at the last window position, which is exactly the one position per window where the teacher-forced level-0 side info is necessarily the zero/unknown padding (the window never contains the token one-beyond-itself). This systematically understated the 2-level model, invalidating comparisons against the 1-D prior. Now scores level-1 at the second-to-last position, where the teacher-forced level-0 code is real. All 488 tests still pass; ruff check/format clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
|
Addressed all 6 findings from the Copilot review (commit f9f496d):
All 488 tests pass, ruff check/format clean, CI green. Merging. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
In-flight research work exploring entropy coding as a swappable second stage on top of the existing frozen codecs (RVQ, SPIHT, Hybrid) — never touches reconstruction quality, only how compactly a codec's own output symbols are packed into bits.
This is not proposing a change to any golden/production configuration yet. It's the framework, tooling, and a set of research findings, in good shape for review before deciding what (if anything) gets promoted to a golden release.
What's included
compressionkit/runtime/entropy_algorithms.py— pluggableEntropyCoder-conforming wrappers:LearnedEntropyCoder(AI prior + arithmetic/rANS backend) and classicalStaticHistogramEntropyCoder/MarkovEntropyCoderbaselines, so "does the AI prior actually help" can be measured by swapping one object.compressionkit/generative/causal_priors.py— causal token-prior architectures (WaveNet/GRU/CNN/DSCNN/hybrid variants) sharing a uniform(B, context_length) -> (B, context_length, vocab_size)interface.compressionkit/generative/spiht_bit_prior.py— a pure-numpy inference path for a neural bit-level prior on SPIHT's bitstream, injected via a new genericsink/sourcedependency-injection seam added tocompressionkit/dsp/spiht.py(SPIHT's tree-traversal algorithm stays entirely neural-agnostic).compressionkit/runtime/prior.py,compressionkit/runtime/two_stage.py— real correctness fixes:predict_log_probs/_get_probsnow use the same windowed-context convention as real incremental/online inference (previously used a full unmasked vectorized shift that didn't match actual context-truncation behavior)._CDF_TOTALvs_WHOLE) that could cause an infinite renormalization loop with confident, real (non-uniform) trained-prior probabilities."rans"backend alongside the existing arithmetic coder.docs/entropy-coding-research.md— the research framework doc: blocks, ready-made experiments, golden-promotion path, and "known lessons" (single-level token bug, LR-schedule/epoch-vs-steps pitfalls, augmentation reuse, validation-split leakage) so these aren't re-derived in a future session.scripts/measure_rvq_entropy.pyadditions,scripts/measure_rvq_entropy_2level.py,scripts/benchmark_rans_vs_arithmetic.py,scripts/build_augmented_ecg_train_cache.py) and corresponding tests (47 new/modified acrosstests/test_causal_priors.py,tests/test_entropy_algorithms.py,tests/test_entropy_benchmark.py,tests/test_spiht_bit_prior.py,tests/test_two_stage.py).Research findings (see
docs/entropy-coding-research.mdfor detail)Test/lint status
uv run pytest -q: 488 passed.uv run ruff check ./uv run ruff format --check .: clean (fixed a real missing-import bug inspiht.pyplus lint/format drift in 6 files as part of this PR).Not in this PR