wai-rs: derive the dequant gain — the last fixture-sourced decoder input #29
2026-07-26
wai-rs: land the f32 synthesis variants the decoder depends on
Some checks failed
smart-byte conformance gate / Validate conformance JSON files (pull_request) Successful in 3s
wai cross-architecture byte-equality / macos-arm64 (reference) (pull_request) Failing after 5s
wai cross-architecture byte-equality / linux-aarch64 (pull_request) Failing after 4s
wai cross-architecture byte-equality / windows-x86_64 (pull_request) Failing after 4s
smart-byte conformance gate / Spec changes require vector-pack updates (pull_request) Successful in 3s
WAI conformance / cargo-class corpora · linux-arm64 (pull_request) Failing after 3s
smart-byte conformance gate / VERSIONS.md and README.md must exist (pull_request) Successful in 3s
WAI conformance / cross-arch byte-equality · linux-arm64 (pull_request) Failing after 3s
joule-code conformance gate / Build + round-trip C1-C4 vectors (pull_request) Failing after 2s
WAI conformance / cargo-class corpora · windows (pull_request) Failing after 9s
jouleclaw public surface must stay documented / Public-surface audit (jouleclaw-rs) (pull_request) Failing after 2s
no-vanity-metrics / Scan committed Markdown for vanity metrics (pull_request) Successful in 3s
WAI conformance / cross-arch byte-equality · windows (pull_request) Failing after 10s
wai cross-architecture byte-equality / linux-x86_64 (pull_request) Failing after 2s
WAI conformance / cargo-class corpora · linux (pull_request) Failing after 3s
WAI conformance / cross-arch byte-equality · linux (pull_request) Failing after 2s
WAI neural decoders end-to-end / synthetic fixtures → Rust decode → reference (±4) (pull_request) Failing after 45s
WAI conformance / cargo-class corpora · macos (pull_request) Successful in 1m24s
WAI conformance / cross-arch byte-equality · macos (pull_request) Failing after 18s
decode_y_f32 / decode_uv_f32 were in the working tree but never committed, while
the whole-decoder commit and jpegai_decode both call them. So the branch as pushed
did not compile — the local builds only passed because the functions were sitting
uncommitted on disk. Caught by a pre-merge build across features.
The change factors the output denormalize out of decode_y / decode_uv into
denorm_q16 + denorm_u8 / denorm_f32, with the pre-denorm planes split into
decode_{y,uv}_planes. The u8 decoders are unchanged in behaviour; the f32 variants
return the same pixels before the final round, which is what the post-filter chain
needs (it works in float and rounds once, at BT.709).
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: a library decode() entry point for the JPEG-AI decoder
The whole-decoder logic lived in the end-to-end test's helper. Extract it into `jpegai_decode`, a public `decode(codestream, weights, norm) -> Rgb` that composes the validated stages — entropy front half, integer value path, synthesis, and the post-filter chain — behind one call. The module is pure compute: the caller passes the model weights (DecoderWeights) and normative tables (NormTables); it reads no files. So the decoder is now a callable API rather than a pipeline that only exists inside a test. The end-to-end test is reduced to loading the weights, calling decode(), and comparing — and gives the identical result, max|Δ| = 1 and 4.446% of samples differing, confirming the extraction preserves behaviour. decode() rejects unsupported geometries up front (returns None): the shipped profile is 256×256 4:4:4 tools_0, luma 160ch / chroma 96ch. The domain and wiring facts the composition depends on — z at 4×4 vs the 16×16 plane, the [0,1] eICCI domain, EFElinear's inverted filters_best2, the shared SOZ stream — are documented at their call sites so the next reader doesn't re-derive them. 107 neural_int tests pass. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: the whole decoder end to end — codestream → RGB vs the reference PNG
Composes every validated stage into one decode of test256.bin and compares the
result to the reference test256_dec.png:
codestream ─┬ luma: z→i_s→σ/mask→SORP→y_q→resi→ψ→MCM→ŷ→LSBS→decode_y
└ chroma: z→i_s→σ/mask→SORS→y_q→resi→ψ→(no MCM)→ŷ→LSBS→decode_uv
YUV → EFElinear → eICCI → EFEnonlinear → LEF → BT.709 → RGB
Result: max|Δ| = 1, 4.45% of samples differ. Every differing sample is off by
exactly one, and the fraction sits at the PTQ floor already measured per plane —
luma ~1.4%, chroma ~3.4%, so ~4.4% across the combined RGB is expected. This is
not bit-equality with the reference PNG and is not meant to be: ψ, the MCM stages
and synthesis are 16-bit PTQ. The contract the integer path keeps is determinism —
byte-identical output on every machine — bought with ~1 LSB of accuracy.
Composing the stages is what earned this test its place. Each stage passed in
isolation, yet the composition surfaced four wiring faults no per-stage test could
see, three of them the exact class this test exists to catch:
* hyper_scale / hyper_decoder take z at the 4×4 hyper resolution, not the 16×16
latent plane — a hardcoded 16 indexed off the end of the weights;
* reconstruct_yhat / LSBS take the 16×16 plane dims, which the function halves
internally — passing 2× double-counted;
* eICCI runs in [0,1]; the chain handed it [0,255], the folded BN and ReLUs
misbehaved, and it scored worse than a no-op (already fixed in the eICCI
commit, re-confirmed composed here);
* EFElinear applies `filters_best2`, NOT `filters_best` — the naming is inverted,
and the wrong group has a ~−16 luma cross-gain that blew U/V to −3882.
wai_jpegai_build_e2e_inputs.py extracts every runtime input — the model weights
for both planes, the normative tables, the LSBS and LEF op-point tables, and the
reference RGB — straight from the shipped assets, so the test is reproducible from
the assets alone rather than from hand-made scratch files. Verified: regenerating
the inputs with it and re-running gives the identical 4.45%.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: validate LEF — the post-filter chain is now complete
Last of the four. Reproduces the reference to 1.5e-5 on Y, against an identity baseline of 1.74, with U and V asserted to pass through untouched (LEF sharpens luma only). LEF keys its sharpening off the raw per-channel scale surface at the latent's 16x16 resolution, for the channel the header names, thresholded against the picture model's tables. Both the channel index and the model id are cross-checked against the decoded header rather than assumed. That sigma dump is also an independent confirmation of the RVS additive derived earlier: subtracting the pre-RVS surface from it leaves exactly [-31, 20, 122] — the derived per-band constants for the GRFS-set branch, arrived at from a different dump than the one they were derived from. All four post-filters now reproduce the reference from header-decoded parameters: EFElinear 1.5e-5 (U, V) identity 2.06 / 1.36 eICCI 6.1e-5 (Y,U,V) identity 6.45 / 2.56 / 2.72 EFEnonlinear 4.6e-5 (V) identity 8.53 LEF 1.5e-5 (Y) identity 1.74 Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: validate EFEnonlinear against its reference dump
Third stage of the post-filter chain, parameters decoded from the tools header: the eight filter weights and the per-tile luma range. Reproduces the reference to 4.6e-5 on V, against an identity baseline of 8.53. This picture enables the filter on V only. Both the header and the dump say so, and the test asserts both — that Y and U pass through byte-identical, and that the header's plane flags agree. A header misparse that silently flipped the plane selection would otherwise look like a pass on the plane it left alone. The decoded luma range is 4.97..253.20, which is exactly the data range of the filter's own input — an independent check that the header offset is right. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: validate EFElinear against its reference dump, from header-decoded params
First stage of the post-filter chain. Parameters come from the tools header decoded out of the codestream — the two chroma plane means and the self/cross gains — not from a fixture. Reproduces the reference EFElinear to 1.5e-5 on both chroma planes, against identity baselines of 2.06 and 1.36. Two things about the header parse are worth pinning down, because getting either wrong produces plausible-looking garbage rather than an error: the post-filter chain starts 2 bits into TON (other tools emit ahead of it), and models_count sets the signalled-index width. With the wrong offset the parse desynchronised into weight vectors of 200+ entries and chroma plane means of 277 — nonsense for 8-bit data, but nothing raised. The applied weights are filters_best2, not filters_best. Also asserts the dumped filters chain end to end (out[i] == in[i+1]), which is what licenses feeding each stage the reference's own input rather than the previous stage's output: a defect then fails in the filter that caused it instead of smearing downstream. Every plane is reported against its identity baseline. These filters are small corrections that each touch only part of the image — EFElinear the two chroma planes, EFEnonlinear only V, LEF only Y — so a bare closeness check would pass on a no-op. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: validate eICCI as a whole filter against the reference dump
IcciPath::forward_plane was already checked against process_comp, but only per plane and only from pre-computed inputs. This runs the whole filter — the two-level Haar transform of all three planes into the 48-channel input, the luma path, the shared chroma trunk with its two heads, the inverse transform — against the reference's own EfficientICCIFilter in/out dump on the real test256 YUV. It reproduces the filter to float epsilon: max|Δ| of 6.1e-5 / 6.1e-5 / 4.6e-5 on Y / U / V, against identity baselines of 6.45 / 2.56 / 2.72. The domain is the part worth recording. eICCI runs on [0,1] data — the reference normalises by 255 before the filter and scales back after. Fed raw 0..255 the folded BatchNorms and ReLUs misbehave and the network scores *worse than doing nothing* (7.12 against a 6.45 identity baseline). That failure looks like a plausible near-miss rather than an error, so every plane is asserted against its identity baseline: the output must sit far closer to the reference than the input does, and a no-op fails instead of passing quietly. Co-Authored-By: Claude Opus 4.8 <[email protected]>
tools: an icci export mode for the eICCI model weights
Extracts the IcciNet weights from the two eICCI wrappers into the JSON the icci calibrations read. Heads and scales come out by name; the trunk convolutions are anonymous in the graph and are assigned by file order — conv_first, then each residual block's 1x3 tap followed by its 3x1 tap. These wrappers are BN-pre-folded, unlike post_filters/model_*.onnx, so there is nothing to fold. The structure is asymmetric in two ways worth stating: the luma path has 2 residual blocks where chroma has 4, and the blocks are factorized into 1x3 and 3x1 taps rather than a single 3x3. Which pair of wrappers applies is signalled per picture — icci_model_signalled_idx indexes a shortlist, which for the 256x256 conformance stream selects Y = model_6 and UV = model_3. Having this as a mode rather than a one-off script means the weights are regenerable; the previous extraction was lost with its scratch directory. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: eICCI — give the V head its own residual scale
The reference signals comp_model.scaleU and comp_model.scaleV separately, and on the shipped chroma model they differ by about 1% (0.011744 vs 0.011875). IcciNet applied a single shared chroma.scale to both heads, which shifts V by more than a pixel LSB. Adds IcciNet.scale_v for the V head; the U head keeps chroma.scale. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: chain the chroma plane — codestream to decoded U,V pixels
The chroma decode now runs end to end alongside luma, off raw bytes with no
per-image fixture:
codestream → z_uv → i_s → sigma/mask/rvs → y_q 24576/24576 bit-exact
→ resi 24576/24576 bit-exact
→ ψ → y_hat → uv_lsbs
→ decode_uv(cat(y_lsbs, uv_lsbs)) U,V[256×256], max|Δ| = 1
Chroma is structurally different from luma, not a reparameterisation:
* No MCM. The reference ships no MCM/ nets under model_uv; the hyper params
are used directly as the means. Added reconstruct_yhat_uv — a plain
up-shuffle and add — rather than flagging reconstruct_yhat, because the two
paths genuinely differ.
* Joint synthesis. decode_uv consumes the concatenated 256-channel latent
cat(y_lsbs, uv_lsbs), so chroma pixels depend on the luma chain too; a luma
regression surfaces in this test as well.
* One SOZ substream carries both planes' hyper-latents back to back through a
single ANS decoder, so the luma pass is decoded first to leave the stream at
the right position for chroma.
The synthesis nets are generalised to read their channel widths from the
weights (decode_hyper, decode_hyper_scale) — chroma is the same topology at
96/384/1536 where luma is 160/640/2560, so nothing may assume the luma count.
grfs and beta_displacement are now DERIVED from the picture header via
parse_picture_header (component [0] luma, [1] chroma) rather than read from the
fixture — the same parser jpegai_header_parser_matches_trace checks against the
reference trace. The luma test is switched over too. That removes the last
per-image scalar inputs from both e2e tests: everything now comes from the
codestream plus normative assets and sink weights.
Pixel error, attributed rather than assumed: chroma differs on 3.38% of samples
(all by exactly 1). Feeding the reference latent into uv synthesis alone gives
2.06%, so the synthesis PTQ dominates and the whole upstream chain adds 1.32%.
Chroma runs looser than luma (1.4%) because its psi is 60.6 dB vs luma's 70.5
and, with no MCM to refine them, that quantization reaches the latent
unattenuated. Bar set at 4% with max|Δ| <= 1 held strictly.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: chain synthesis — codestream to decoded luma pixels
The pure-integer path now runs end to end, raw bytes to pixels, with no
per-image fixture:
codestream → y_q → resi bit-exact
→ ψ → y_hat → y_lsbs
→ decode_y Y[256×256], max|Δ| = 1
Every differing sample is off by exactly 1. About 1.4% of samples differ, which
is looser than the reference decoder's own 0.1% bar, so the test attributes the
gap rather than just widening the threshold: feeding the *reference* latent into
synthesis alone already yields 1.202%, while the full chain yields 1.398%. The
16-bit PTQ of the synthesis net is therefore the dominant term, and everything
upstream — the entropy decode, dequant, hyper decoder, four MCM stages and LSBS
— contributes about 0.2%.
That is the honest cost of an integer-throughout decode. The reference
comparison runs float nets for ψ, MCM and synthesis; this path does not, and
trades ~1.3% of samples at 1 LSB for determinism: byte-identical output on every
machine, which is the property the contract is about.
The bar is set at 2% with max|Δ| <= 1 held strictly, and the isolation
measurement is printed on every run so a regression in the chain is
distinguishable from the fixed quantization floor.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: fix the LSBS block statistic, and chain LSBS onto the value path
lsbs_refine computed its own block statistic from a slice it called `i_s`,
summing magnitudes. Both parts were wrong. The summed quantity is the signed
gain-adjusted scale (i_s + gain_per_ch) — which genuinely goes negative, min
-687 on this image — so summing magnitudes reproduces only 1536/40960 of the
reference's block statistics on real luma data.
This is the same error just fixed in rvs_additive, in a third copy of the same
logic. It matters more here: the RVS additive buckets `likely` into one of four
bands and so absorbed the error, while LSBS indexes a 4352-entry table with it
directly, and reads the wrong coefficients.
So the statistic is no longer recomputed here. lsbs_refine now takes `likely`,
and every consumer — rvs_additive, dequantize_resi, lsbs_refine — is driven by
the single rvs_likely derivation. Verified against the reference: the LSBS and
dequant `likely` dumps are byte-identical, which is what licenses sharing it.
Chained into jpegai_e2e_frontend (y_hat → LSBS → synthesis input) and the
standalone calibration rebuilt on the reference's own pp{0,1} dumps, so it now
covers both planes and tests the formula in isolation from upstream PTQ error:
bit-exact, 40960/40960 luma and 24576/24576 chroma, max|Δ| = 0.
Both tests also assert against the IDENTITY baseline, because that is what makes
them mean anything. LSBS is a small correction (max|Δ| ≈ 0.079), so passing y_hat
through unrefined already scores 51 dB against the refined reference — the 40 dB
bar I first wrote would have passed a no-op. The bar is now 65 dB, above identity
and below the 74.5 dB actually achieved, and the isolated test requires the
identity error to exceed the refinement error by 100x.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: chain the value path — codestream to the synthesis-input latent
Extends the chain through reconstruction, so the whole front half now runs off
raw bytes in one pass:
codestream → z → i_s → sigma/mask/rvs → y_q 40960/40960 bit-exact
→ resi 40960/40960 bit-exact
→ ψ (hyper decoder) SNR 70.48 dB
→ y_hat (SGMM + 4 MCM stages) SNR 74.54 dB
Two things worth noting beyond the wiring.
The full dequant — RVS composed with the gain unit — reproduces the reference
residual bit-exactly, 40960/40960. That confirms the earlier decision to apply
the two stages sequentially rather than fold them into one per-position gain:
the fused form was algebraically identical but drifted by a last-bit f32
rounding, and would have shown up here as an inexact result.
The reconstruction is autoregressive across four MCM stages, each conditioned on
the sub-lattices reconstructed before it, so an error in stage i corrupts every
stage after it. Together with the mANS state threading across all 160 channel
calls upstream, that makes this a single end-to-end assertion: essentially
nothing in the chain can be wrong while the result still lands at 75 dB.
Integer stages are asserted bit-exact; the PTQ nets are asserted on
signal-to-error, since bit-equality with a float net is not achievable by
construction. Both PTQ bars are stated with their do-nothing baselines.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: chain ψ onto the front end, and guard the Q16 input contract
Extends the front-end chain with the hyper decoder: the same z that drives the scale net also produces ψ, the MCM entropy context. SNR 70.48 dB against the float reference, off the codestream. The stage is a PTQ net, so the check is signal-to-error rather than bit-equality, and it is reported against the do-nothing baseline — because this particular stage fails in a way that is easy to mistake for working. decode_hyper requires Q16 input. Entropy-decoded symbols are small plain integers (about -3..4), so feeding them unscaled is an easy mistake, and it does not fail loudly: the first conv comes out 2^16 too small, ReLU6 clamps it to zero, and the result collapses to conv4's bias — which still looks like a smooth surface of roughly the right magnitude. Measured, that mistake lands at rmse 0.074 against a zero-baseline of 0.079: 0.7 dB better than emitting nothing, while a correctly lifted input scores ~70 dB. I made exactly that mistake while wiring this up and briefly concluded the decoder was broken. Three guards so the next reader doesn't: * q16_from_symbols() names the lift at the call site * decode_hyper's doc states the hazard with both measurements * the chain test asserts SNR > 45 dB, which the unscaled path cannot reach Also a metric note, since it is what made the failure hard to see: PSNR with peak=255 on a surface whose range is +-1.2 inflates by ~40 dB, enough to make a collapsed-to-bias output read as a passing number. Prefer SNR, and quote the do-nothing baseline beside any dB figure. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: the decoder front end end-to-end, off raw bytes with no per-image fixture
Chains the validated stages into one pass:
test256.bin → SOZ → mANS fac → z
→ decode_hyper_scale → i_s
→ rvs_likely / rvs_additive / skip_mask / sigma_index
→ SORP → mANS sgm → y_q
→ rvs_dequant → resi
Both ends byte-exact on real data: 40960/40960 symbols and 40960/40960
residual, with 703 coded positions derived rather than supplied.
The stages already had unit calibrations, so what this adds is the absence of
per-image input. The σ distribution indices, the coded-position mask, the RVS
band and the dequant multiplier are all derived from the bitstream, and the mANS
state threads across all 160 channel calls — so an error anywhere upstream
desynchronizes the entropy decoder and the symbols come out wrong. That makes
the symbol check a single end-to-end assertion rather than 160 independent ones,
and it is the property the per-stage tests cannot establish on their own.
Inputs are the codestream plus normative assets only: the factorized CDFs, the
me-tANS PMF/bounds tables, the hyper-scale net, gain_unit_mlog.csv.
To get there, build_decode_transitions + cdf_from_pmf are ported to the integer
mANS module, so the entropy stage builds its transition tables from the
normative PMF instead of loading a precomputed fixture — and no longer reaches
into the float/ONNX-side module to do it.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: fix the RVS block statistic — sum the gain-adjusted scale, not |i_s|
`rvs_additive` summed `|i_s|` per 8×8 block. The quantity the reference sums is the gain-adjusted scale, `i_s + gain_per_ch[c]`, with no magnitude. This passed its calibration anyway, and would have kept passing: the two sums differ (12416/40960 agree on luma, 24320/24576 on chroma) but still select the same RVS band for every block in test256, so the additive came out right for the wrong reason. A single-image byte-exact result is not evidence the statistic is correct — only that it is correct enough to land in the same bucket here. The dequant is what distinguishes them: it indexes its Q16 multiplier table by this value directly rather than bucketing it, and is bit-exact only with the gain included — 40960/40960 luma and 24576/24576 chroma with `i_s + gain`, 12416 and 24320 without. Extracted as `rvs_likely` and shared by both callers so they cannot drift apart, and the dequant calibration now derives it from the scale surface instead of reading it from the fixture, so the derivation itself is under test. Co-Authored-By: Claude Opus 4.8 <[email protected]>
wai-rs: derive the dequant gain — the last fixture-sourced decoder input
The per-position dequant gain was the one input still read from a dumped
fixture. It is now built from normative assets alone, closing the derivation
chain: every input the decoder needs comes from the codestream plus the
shipped model assets.
Two stages, applied in the reference's order:
RVS out = y_q · round(2^23 / scale[band]) / 2^16
gain unit resi = out / (scaler_vec[c] + 1e-9)
with scale[i] = round_ties_even(base[i] · (grfs_flag ? 1.06 : 0.98)) over the
normative base list [115, 125, 146, 160], and
scaler_vec[c] = round(exp((gain_vector_log[c] + beta) · log_k / 2^7) · 2^10) / 2^10
where gain_vector_log is the normative gain_unit_mlog.csv and beta comes from
the picture header. scaler_precision is gain_vector_precision +
beta_displacement_precision = 5 + 5 = 10.
The same scale list also reproduces RVS_ADDITIVE exactly via
round(ln(scale/128) / (log_k / 2^7)) — those constants were previously carried
as a table, and are now cross-checked against their derivation.
Verified on real test256.bin data: the RVS stage is bit-exact against the
reference on both planes (luma 40960/40960, chroma 24576/24576), and the
derived tables match the reference's own buildTables() output for all three
cwg variants.
Three details worth recording, each of which silently corrupts the result:
* The scale rounding is half-to-even. 125·1.06 is exactly 132.5 and 125·0.98
exactly 122.5 in binary64; half-away-from-zero yields 133/123 and selects
the wrong multiplier for that band. Uses round_ties_even.
* The reference's dequantize_resi has two paths — integer (x·t)>>16 and float
(x·t)/2^16 — selected by a device comparison rather than by the bitstream.
They disagree for negative x (floor vs exact). The float path is the one it
executes; that is what this reproduces.
* Folding the two stages into a single per-position gain is algebraically
identical but differs in the last bit of the f32 result, so they are applied
sequentially.
Note the two log_k constants differ and are easy to transpose: the gain unit
uses (ln 54.82 − ln 0.11)/31, the residual-variance scale uses 54.85. The
additive table is built with the 54.82 one.
Co-Authored-By: Claude Opus 4.8 <[email protected]>