Transaction-Science/open-standards
Watch
1
0
Fork
You've already forked open-standards
 0
Code Issues Pull requests Projects Releases 15 Packages Wiki Activity Actions

wai: D2c O8 + M7b — JPEG AI's networks written in jpegai-int/1 and run through wai-ops/1, goldens unchanged #315

Merged
dcharlot merged 15 commits from wai/d2c-m7b into main 2026-10-18 12:35:32 -04:00
Conversation 0 Commits 15 Files changed 71 +3098 -323
71 changed files with 3098 additions and 323 deletions
Download patch file Download diff file Expand all files Collapse all files

7
.github/workflows/wai-conformance.yml vendored
View file

@ -361,16 +361,19 @@ jobs:
# jpegai_ops.rs holds the JPEG AI scale path (integer-operators §13.6.1) to its committed
# lowered graphs and certificate here; its cases on the jpegai/v1 bundle skip here and
# run in wai-jpegai-e2e.yml, which fetches the bundle and requires them.
# jpegai_nets_ops.rs does the same for its other integer networks (§13.6.2): the
# committed synthetic networks' lowered graphs and certificate, and every kernel
# policy on them, here; its jpegai/v1 cases there.
# int_video_ops.rs and wir1_ops.rs hold WIV1 and WIR1, which decode through the
# operator set's stage operators with no graph (integer-operators §13.5), to their
# stage forms, goldens and refusals; `i32_forms` builds E9's 32-bit form, which they
# and int_ops_differential.rs hold to the exact one.
- name: Fast integer kernels ≡ reference kernels (overflow checks on)
working-directory: wai/wai-rs
run: cargo test --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int,i32_forms --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_hyper_ops --test int_mlicpp --test int_mlicpp_preflight --test int_mlicpp_ops --test int_world_ops --test int_world_memory --test jpegai_ops --test int_video_ops --test wir1_ops --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_set_cache --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration
run: cargo test --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int,i32_forms --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_hyper_ops --test int_mlicpp --test int_mlicpp_preflight --test int_mlicpp_ops --test int_world_ops --test int_world_memory --test jpegai_ops --test jpegai_nets_ops --test int_video_ops --test wir1_ops --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_set_cache --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration
- name: Fast integer kernels ≡ reference kernels (release, overflow checks off)
working-directory: wai/wai-rs
run: cargo test --release --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int,i32_forms --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_hyper_ops --test int_mlicpp --test int_mlicpp_preflight --test int_mlicpp_ops --test int_world_ops --test int_world_memory --test jpegai_ops --test int_video_ops --test wir1_ops --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_set_cache --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test manifest_parse_memory --test manifest_reader_equivalence --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration
run: cargo test --release --no-default-features --features neural_int,video_int,audio_int,world_int,splat_int,generate,score,staged,stage_measure,legacy_int,i32_forms --lib --test int_kernels_vectors --test int_kernels_geometry --test int_synth_ops --test int_ops_differential --test int_splat_ops --test int_audio_ops --test int_hyper_ops --test int_mlicpp --test int_mlicpp_preflight --test int_mlicpp_ops --test int_world_ops --test int_world_memory --test jpegai_ops --test jpegai_nets_ops --test int_video_ops --test wir1_ops --test int_audio --test int_splat --test generate --test cli_generate_score --test int_payload --test int_payload_preflight --test int_payload_cost --test int_set_cache --test int_video_enc_memory --test cli_stage --test staged_resources --test stage_measure --test stage_measure_bounds --test manifest_parse_memory --test manifest_reader_equivalence --test int_video --test int_video_halfpel --test int_video_keyframe --test int_world --test int_world2 --test jpegai_bitstream_calibration
# The integer decode path builds under each of its features alone, so a sink can take
# one medium without the others, and for wasm32 (a browser sink); `int_payload` arms a
# medium only under its feature, and none needs `neural`. So does the staged-measurement

58
.github/workflows/wai-jpegai-e2e.yml vendored
View file

@ -24,10 +24,12 @@ name: WAI JPEG-AI whole decode (pinned inputs, required)
# - two decodes in one process are byte-identical and the reconstruction binding
# verifies the second and rejects a one-byte change;
# - the bundle's bounds certificate (wai/extensions/integer-operators.md §8.5,
# §13.6.1: the two scale graphs) is byte for byte the one committed under the
# set's digest in wai/ops-conformance/certificates/, which is the listing digest
# the set-form pin below covers, from both certifiers (wai-verify and the
# independent tools/wai_ops_check.py), and its verdict meets acc32;
# §13.6: the two scale graphs and the four graphs of its other integer
# networks) is byte for byte the one committed under the set's digest in
# wai/ops-conformance/certificates/, which is the listing digest the set-form
# pin below covers, from both certifiers (wai-verify and the independent
# tools/wai_ops_check.py); its verdict meets acc64, and the scale path's graphs
# meet acc32 (§11, role_profiles);
# - the decoded pixels hash to WAI_JPEGAI_EXPECT_SHA256 below: the value measured
# from these inputs on an Apple Silicon Mac (aarch64 macOS, native) and
# reproduced by an x86_64 Linux build (glibc, run under QEMU user-mode
@ -111,6 +113,7 @@ on:
- 'wai/wai-rs/tests/jpegai_reconstruction_binding.rs'
- 'wai/wai-rs/tests/jpegai_cross_platform.rs'
- 'wai/wai-rs/tests/jpegai_ops.rs'
- 'wai/wai-rs/tests/jpegai_nets_ops.rs'
- 'wai/wai-rs/tests/jpegai_memory.rs'
- 'wai/wai-rs/src/int_set_cache.rs'
- 'wai/wai-rs/src/int_ops.rs'
@ -151,6 +154,7 @@ on:
- 'wai/wai-rs/tests/jpegai_reconstruction_binding.rs'
- 'wai/wai-rs/tests/jpegai_cross_platform.rs'
- 'wai/wai-rs/tests/jpegai_ops.rs'
- 'wai/wai-rs/tests/jpegai_nets_ops.rs'
- 'wai/wai-rs/tests/jpegai_memory.rs'
- 'wai/wai-rs/src/int_set_cache.rs'
- 'wai/wai-rs/src/int_ops.rs'
@ -203,7 +207,7 @@ jobs:
if: needs.relevant.outputs.run == 'true'
name: fetch pinned JPEG-AI inputs → whole decode + pixel hash, required
runs-on: linux
timeout-minutes: 30
timeout-minutes: 60
env:
WAI_REQUIRE_SOTA_VECTORS: '1'
WAI_JPEGAI_DIR: ${{ github.workspace }}/jpegai_e2e
@ -229,7 +233,7 @@ jobs:
# bundle is not committed, so its certificate is, under the set's digest (the
# listing of norm.json and weights/*, as the set-form pin covers it), and
# listed in wai/ops-conformance/certificates/external.json.
- name: The bundle's bounds certificate is the committed one, and meets acc32
- name: The bundle's bounds certificate is the committed one; it meets acc64, and its scale path acc32
shell: bash
working-directory: wai
run: |
@ -240,7 +244,8 @@ jobs:
for n in hs_y hs_uv hyp_y hyp_uv synth_y synth_uv mcm0_y mcm1_y mcm2_y mcm3_y icci; do
M+=(--member "weights/$n.json=$D/$n.json")
done
wai-verify/target/release/wai-verify bounds --capability wai.image.jpegai "${M[@]}" --profile acc32 \
wai-verify/target/release/wai-verify bounds --capability wai.image.jpegai "${M[@]}" \
--profile acc64 --role-profile hs_y=acc32 --role-profile hs_uv=acc32 \
--out "${RUNNER_TEMP:-/tmp}/jpegai-bounds.json" \
--check ops-conformance/certificates/wai.image.jpegai/3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37.json
# The same certificate from the independent certifier (Python standard library
@ -248,7 +253,7 @@ jobs:
- uses: actions/setup-python@v5
with:
python-version: '3.11'
- name: The independent certifier writes the same certificate, and it meets acc32
- name: The independent certifier writes the same certificate; it meets acc64, and its scale path acc32
shell: bash
working-directory: wai
run: |
@ -259,7 +264,7 @@ jobs:
M+=(--member "weights/$n.json=$D/$n.json")
done
python3 -I tools/wai_ops_check.py certify --capability wai.image.jpegai --lowering jpegai-int/1 "${M[@]}" \
--profile acc32 --out "${RUNNER_TEMP:-/tmp}/jpegai-bounds-py.json" \
--profile acc64 --role-profile hs_y=acc32 --role-profile hs_uv=acc32 --out "${RUNNER_TEMP:-/tmp}/jpegai-bounds-py.json" \
--check ops-conformance/certificates/wai.image.jpegai/3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37.json
# The operator set's lowering of the bundle's scale path (integer-operators
# §13.6.1), run through the interpreter with every value held to its certified
@ -340,3 +345,38 @@ jobs:
echo "::error::a JPEG-AI test skipped instead of checking its input"
exit 1
fi
# The other integer networks through the integer operator set (integer-operators
# §13.6.2, lowered by jpegai-int/1 to rec_y, rec_uv, syn_y and syn_uv), on this
# bundle: the lowering of the networks as the decoder parses them is the lowering of
# their bytes and the committed certificate's graphs, under acc64; each network
# through its graph gives the direct network's values on the test codestream; every
# value each operator forms there lies within its certified interval (the measured
# accumulators of §13.6.2); and every kernel policy runs the graphs alike.
- name: The hyper decoders, context stages and synthesis through the integer operator set, on the pinned bundle (missing input = failure)
shell: bash
working-directory: wai/wai-rs
run: |
set -o pipefail
log="${RUNNER_TEMP:-/tmp}/jpegai-nets-ops.log"
[ "${WAI_REQUIRE_SOTA_VECTORS:-}" = 1 ] \
|| { echo "::error::WAI_REQUIRE_SOTA_VECTORS must be 1 in this job"; exit 1; }
export WAI_JPEGAI_CODESTREAM="$WAI_JPEGAI_DIR/test256.bin"
export WAI_JPEGAI_WEIGHTS_DIR="$WAI_JPEGAI_DIR"
export WAI_JPEGAI_NORM="$WAI_JPEGAI_DIR/norm.json"
cargo test --release --no-default-features --features neural_int,legacy_int --test jpegai_nets_ops \
-- --show-output 2>&1 | tee "$log"
for t in the_committed_bundle_lowers_to_its_committed_networks \
the_pinned_bundle_lowers_its_networks_to_its_certificate \
the_networks_decode_through_the_lowering_as_directly \
the_whole_decode_through_the_networks_is_the_pinned_decode \
the_networks_keep_to_their_certificate_on_the_test_codestream \
every_load_check_of_the_networks_refuses_model_mismatch \
a_value_past_its_range_is_refused_field_range \
every_kernel_policy_runs_the_networks_alike; do
grep -q "^test $t \.\.\. ok$" "$log" || { echo "::error::$t did not run and pass"; exit 1; }
done
if grep -nE 'SKIPPED|NOT VERIFIED' "$log"; then
echo "::error::a JPEG-AI test skipped instead of checking its input"
exit 1
fi

4
.github/workflows/wai-ops-differential.yml vendored
View file

@ -12,7 +12,9 @@ name: wai-ops differential (nightly)
# integer renderer, against its direct renderer with the direct guard, and
# wai.image.jpegai's scale path and entropy decodes, against its direct computation,
# which panics (or, without overflow checks, may compute on) where the operator set
# refuses a stream; and WIV1 (wai.video.int_motion, the P-frames of wai.video.int_hyper)
# refuses a stream, and its hyper decoders, context stages and synthesis (the graphs of
# integer-operators §13.6.2), against its direct networks, which compute on where the
# operator set refuses a value past its i32 checks; and WIV1 (wai.video.int_motion, the P-frames of wai.video.int_hyper)
# and WIR1 (the four int_refine capabilities), which decode through the stage operators
# alone and make no planned break.
#

43
wai/CHANGELOG.md
View file

@ -6,6 +6,48 @@ wire format may change and such changes are recorded here.
## v0.4.0 — current
### `wai.image.jpegai`'s hyper decoders, context stages and synthesis through the integer operator set
- **§13.6.2 is normative.** [integer-operators](extensions/integer-operators.md) §13.6.2
writes the topology of the capability's other integer networks as four `WHG1` graphs,
lowered by `jpegai-int/1` from the bundle's eight weight files in §13.0's canonical
form: `rec_y` and `rec_uv` (each plane's hyper decoder, and for luma its four context
stages, from the hyper-latent and the dequantised residual to the reconstructed
latent) and `syn_y` and `syn_uv` (the synthesis). §1's claim widens to what these
networks give for their inputs. The independent checker (`tools/wai_ops_check.py`),
written from the text alone, reproduces every lowered graph and certificate byte for
byte.
- **Ranges.** Every activation an operator can grow, and every input but the
hyper-latent, is checked within i32 (`A32`), so every sum is bounded by the weights.
On the test codestream inputs stay within 21 bits, activations within 24, and the
largest sum is 37 bits against certified bounds of 42 to 60: `acc64` holds for every
sum, and no wider accumulator is needed (§13.6.2, *Accumulators, measured*).
- **Profiles.** The `ops` member is `{ wai-ops/1, acc64, full, jpegai-int/1,
role_profiles: { hs_y: acc32, hs_uv: acc32 } }`: the set meets `acc64`, and its scale
path still `acc32` (§11's new `role_profiles`; both certifiers take
`--role-profile`). The certificate of `jpegai/v1` (`3e243dc1…`) holds the six graphs.
- **The switch.** `jpegai_decode::{decode, decode_with_stages}` run the hyper decoders,
context stages and synthesis through the operator set (`jpegai_ops::Networks`,
`jpegai_decode::load_networks`); `decode_loaded` takes the scale path and the networks
loaded once, with the eICCI net. Every committed and pinned output is unchanged: the
`Synthesis` stage hash and every other stage hash of `tests/jpegai_cross_platform.rs`,
and the decoded RGB SHA-256 that `wai-jpegai-e2e.yml` pins. The direct networks run
only in `jpegai_decode::legacy`, test-only (`legacy_int`). The neural sink lowers and
loads the networks from the members' bytes through the store of loaded sets.
- **No loader panics.** `Weights::from_json`, `IntNet::from_json` and
`IcciNet::from_json` return a `Result` (`jpegai_synthesis::BadMember`); a mistyped
member is refused `model_mismatch` where it panicked, and the sink reports that
refusal rather than decoding through the float path.
- **Planned breaks** (integer-operators §15.3 item 4, SPEC §8): a weight file that
breaks §13.6.2's member rules or shapes, or networks whose certificate does not meet
`acc64` (`model_mismatch`); a dequantised residual, a latent or an activation of the
networks past i32 as a Q16 value (`field_range`), where the decoder computed on. No
committed bundle or codestream reaches either.
- **CI.** `tests/jpegai_nets_ops.rs` runs in both fast-kernel steps of
`wai-conformance.yml` on the committed synthetic networks, and in `wai-jpegai-e2e.yml`
on `jpegai/v1`, each pinned case required by name; the differential gains the
networks' half, per push and nightly.
### wai-quantum 0.6.3: the wide stabilizer tableau, matrix product states at any angle, and fewer passes for the state vector, noisy shots and sparse Pauli sums
Every pinned digest is unchanged, and the state vector and the sparse-Pauli engine give the
@ -170,6 +212,7 @@ same bits as before.
[integer-operators](extensions/integer-operators.md) §13.5 alone, and register
`{ wai-ops/1, stage: rans64 t4097/lenient/canonical }`: with no parameter set, the
member has no profile, subset or lowering (§11; the registry schema now requires
`set`, those three together, and a lowering or a stage).
- **The stage operators the drivers need.** `wai-ops`' E1 decodes a symbol against a
table in whatever width a payload's parser keeps it (`Rans::symbol_in`, and
`table_ok` reads its entries likewise), and E10, E11 and E12 gain per-element forms

22
wai/SPEC.md
View file

@ -697,7 +697,7 @@ registration and makes no conforming envelope non-conforming (§8).
| `wai.neural.int_synth` | integer fixed-point synthesis (all-i64 transposed-conv, no float/ONNX) | bit-identical on every machine — soundly backs `intent = replicate` | payload and set: [integer-payloads](extensions/integer-payloads.md) §13; decoded through its lowering, [integer-operators](extensions/integer-operators.md) §13.1 (`int_synth/1`); `int_synth.rs` (integer weights sink-supplied; the reference tools make them by training in float, then quantising after training with per-layer shifts calibrated on measured activations, `tools/wai_int_synth_train.py`. Of the committed vectors only `photo` comes from a trained model, `tools/wai_int_synth_real.py`; `gradient`, `bloom`, `facade` and the cross-architecture golden use seeded integer weights) |
| `wai.neural.int_hyper` | integer fixed-point **conditional** codec — whole bitstream→image path integer (integer hyper-decoder σ → integer-σ entropy → integer synthesis) | the integer-entropy contract fully assembled: a PTQ of mbt2018_mean, bit-identical on every machine, whose integer decode CI holds within 1 dB of the same model's float decode on a pinned photograph | payload `WIH1`; `int_hyper_synth.rs` (tables sink-supplied). Payload: [integer-payloads](extensions/integer-payloads.md) §4; its networks decoded through their lowering, [integer-operators](extensions/integer-operators.md) §13.2 (`whm1/1`) |
| `wai.neural.int_mlicpp` | integer fixed-point **MLIC++ entropy topology** — `K` channel slices of the latent, each decoded as a checkerboard anchor pass then a non-anchor pass, with an entropy network over concat[latent decoded so far, hyper] re-run before every pass to give σ‖μ; hyper-decoder, entropy network and synthesis all integer | the channel-autoregressive × checkerboard entropy model at **decode-equivalence**: the encoder and every decoder compute the same σ, so the same distribution for every element, at any thread count; bit-identical on every machine. The committed conformance vector has synthetic weights, so it demonstrates decode-equivalence and states no quality figure | `int_mlicpp.rs` (one-file model bundle sink-supplied; `tools/wai_int_mlicpp_export.py`). Payload and bundle: [integer-payloads](extensions/integer-payloads.md) §14; its networks decoded through their lowering, [integer-operators](extensions/integer-operators.md) §13.4 (`wmb1/1`) |
| `wai.image.jpegai` | JPEG AI (ISO/IEC 6048-1:2025) codestream → RGB — integer entropy front end (mANS, with an integer scale-index derivation) + 16-bit-PTQ integer hyper-decoder, context-model and synthesis networks + dequantisation, latent refinement, post-filters (the eICCI CNN among them) and colour conversion in IEEE-754 float with a fixed evaluation order. Scope: one configuration — 256×256 4:4:4, `model_id` 0 (tools_0), `beta_displacement_log` 0, the high-operation-point synthesis, the eICCI model pair the weight bundle carries, and the other header conditions in [determinism-tiers](extensions/determinism-tiers.md) §5; both decode paths refuse any other codestream. Within it, a sink without the integer weight bundle, or in a non-default floating-point environment, takes a float ONNX decode instead; a refusal of the exact decode (`model_mismatch` or `stream_exhausted`, [integer-operators](extensions/integer-operators.md) §13.6.1) is reported, not replaced by the float decode. Only checkpoint constants are sink-supplied | **entropy-consistency** (registry class `NeuralEntropyExact`): the entropy decode, including the scale-index derivation, is integer on both decode paths, so the decoded symbols reproduce on every machine. The exact decoder also pins one reconstruction to RGB, reproducible across machines under the build and floating-point conditions in [determinism-tiers](extensions/determinism-tiers.md) §5–§6 (it checks the floating-point environment before decoding) — a property of that decoder, not of the capability, so `wai.image.jpegai` does not back `replicate` (§7). On one reference codestream the exact decoder stays within max\|Δ\| ≤ 2 of the reference decoder (recorded max\|Δ\| = 1 on ~4.4% of samples); evidence in [determinism-tiers](extensions/determinism-tiers.md) §6. **Downgraded in v0.4.0:** this row previously claimed decode-equivalence, and the registry classed the capability `StandardDefined` | `jpegai_decode.rs` (+ `jpegai_{bitstream,mans,pipeline,synthesis,icci}.rs`); its scale path and entropy decodes decoded through their lowering, [integer-operators](extensions/integer-operators.md) §13.6.1 (`jpegai-int/1`, `jpegai_ops.rs`) |
| `wai.image.jpegai` | JPEG AI (ISO/IEC 6048-1:2025) codestream → RGB — integer entropy front end (mANS, with an integer scale-index derivation) + 16-bit-PTQ integer hyper-decoder, context-model and synthesis networks + dequantisation, latent refinement, post-filters (the eICCI CNN among them) and colour conversion in IEEE-754 float with a fixed evaluation order. Scope: one configuration — 256×256 4:4:4, `model_id` 0 (tools_0), `beta_displacement_log` 0, the high-operation-point synthesis, the eICCI model pair the weight bundle carries, and the other header conditions in [determinism-tiers](extensions/determinism-tiers.md) §5; both decode paths refuse any other codestream. Within it, a sink without the integer weight bundle, or in a non-default floating-point environment, takes a float ONNX decode instead; a refusal of the exact decode (`model_mismatch` or `stream_exhausted`, [integer-operators](extensions/integer-operators.md) §13.6.1, or `field_range`, §13.6.2) is reported, not replaced by the float decode. Only checkpoint constants are sink-supplied | **entropy-consistency** (registry class `NeuralEntropyExact`): the entropy decode, including the scale-index derivation, is integer on both decode paths, so the decoded symbols reproduce on every machine. The exact decoder also pins one reconstruction to RGB, reproducible across machines under the build and floating-point conditions in [determinism-tiers](extensions/determinism-tiers.md) §5–§6 (it checks the floating-point environment before decoding) — a property of that decoder, not of the capability, so `wai.image.jpegai` does not back `replicate` (§7). On one reference codestream the exact decoder stays within max\|Δ\| ≤ 2 of the reference decoder (recorded max\|Δ\| = 1 on ~4.4% of samples); evidence in [determinism-tiers](extensions/determinism-tiers.md) §6. **Downgraded in v0.4.0:** this row previously claimed decode-equivalence, and the registry classed the capability `StandardDefined` | `jpegai_decode.rs` (+ `jpegai_{bitstream,mans,pipeline,synthesis,icci}.rs`); its integer stages decoded through their lowering `jpegai-int/1` ([integer-operators](extensions/integer-operators.md) §13.6, `jpegai_ops.rs`): the scale path and entropy decodes (§13.6.1), and the hyper decoders, context stages and synthesis (§13.6.2) |
| `wai.generate.spectral` | deterministic integer GENERATIVE synthesis (`intent=create`) — prompt+seed → integer-LUT spectral composition, no float/ONNX | byte-identical on every machine, pinned to (model, prompt, seed, steps); the create-side of "conformance is a hash" (generation-equivalence). Payload: the instruction object of form v0, which names the model by its SHA-256 and lists every size a sink makes, each output checked by its SHA-256 (*Payload of `wai.generate.spectral`*, below) | `generate.rs` (small pinned model bundle) |
| `wai.neural.elic` | checkerboard context model (hyperprior + 2-pass spatial context) | the checkerboard-context entropy class (arXiv 2203.10886) | `neural/context.rs` |
| `wai.neural.mlicpp` | MLIC++ multi-reference entropy (channel-AR slices + checkerboard + global attention) | the multi-reference entropy model (arXiv 2307.15421) | research code, `neural/mlicpp.rs` |
@ -1616,11 +1616,15 @@ registry as `codecs::capability_replicate_class(cap) -> ReplicateClass`:
[determinism-tiers](extensions/determinism-tiers.md) §5. This class is the
neural rule applied to it, not a statement about its standard's conformance on
the reconstruction, which this revision has not read; see *`wai.image.jpegai`
and the standard classes*, below.) The stages the class covers, its scale path
and its two entropy decodes, decode through the integer operator set, named in its
`ops` member ([integer-operators](extensions/integer-operators.md) §11, §13.6.1):
the lowering `jpegai-int/1` and the profile `acc32`, which every conforming
bundle's scale graphs meet; a bundle whose certificate does not is refused.
and the standard classes*, below.) Its integer stages decode through the integer
operator set, named in its `ops` member
([integer-operators](extensions/integer-operators.md) §11, §13.6): the lowering
`jpegai-int/1`, the profile `acc64`, which every graph of a conforming bundle
meets, and the role profile `acc32` for the scale graphs `hs_y` and `hs_uv`, the
stages the class covers with its two entropy decodes; a bundle whose certificate
does not meet them is refused. Its hyper decoders, context stages and synthesis
(§13.6.2) are integer too, from inputs the float stages give, and lie outside
the class's claim.
- **`GenerativeDeterministic`** — the *create*-side of byte-exact conformance.
A generator (`intent = create`) whose entire pipeline is integer/fixed-point
and pinned to `(model, prompt, seed, steps)`, so the **generated** content is
@ -2536,6 +2540,12 @@ envelope that uses any string a later revision registers.
and a latent or hyper-latent substream that ends before the decode it carries,
which panicked (`stream_exhausted`). A sink holding the integer weight bundle
reports the exact decode's refusal rather than decoding through the float path.
Its exact decoder's hyper decoders, context stages and synthesis (§13.6.2): a
weight file that breaks §13.6.2's member rules, which the loaders panicked on, or
whose blocks do not have the shapes the topology names, or networks whose
certificate does not meet `acc64` (`model_mismatch`); and a dequantised residual,
a latent or an activation of those networks that leaves `i32` as a Q16 value,
where the decoder computed on (`field_range`).
- `wai.video.int_motion`, the P-frames of `wai.video.int_hyper`, and the four
`int_refine` capabilities: none. They decode through the stage operators alone
(integer-operators §13.5), with the output and the refusal of the decoders they

7
wai/extensions/determinism-tiers.md
View file

@ -439,7 +439,12 @@ entropy decodes, run through the integer operator set, by the lowering `jpegai-i
([integer-operators](integer-operators.md) §13.6.1) at the profile `acc32`. Their
refusals have codes: `model_mismatch` for a bundle that fails a load check of §13.6.1
or whose scale graphs do not meet `acc32`, and `stream_exhausted` for a latent or
hyper-latent substream that ends before the decode it carries. A sink that holds the
hyper-latent substream that ends before the decode it carries. The exact decoder runs
its hyper decoders, context stages and synthesis through the same operator set too, by
`jpegai-int/1` (§13.6.2) at the profile `acc64`: those graphs' outputs lie outside the
claim, and their refusals are `model_mismatch` for a weight file that fails §13.6.2's
load checks or networks that do not meet `acc64`, and `field_range` for a dequantised
residual, a latent or an activation that leaves i32. A sink that holds the
integer weight bundle and refuses the exact decode reports that refusal, with its code;
it does not then decode the codestream through the float path. Only a sink without the
bundle, or a thread or target that cannot run the exact decoder's float stages, takes

130
wai/extensions/integer-operators.md
View file

@ -12,9 +12,9 @@
> - the canonical lowerings of the existing integer network forms to `WHG1`;
> - the conformance that an operator implementation, a parameter set, a sink and a certifier must each pass.
>
> **What a reader can build from this text.** With the payload extensions it cites, an implementation of this text decodes `wai.video.int_motion`, `wai.audio.int_codec`, `wai.splat.int_codec`, `wai.neural.int_synth`, `wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.neural.int_mlicpp` and `wai.{image,video,audio,splat}.int_refine`, and reproduces the symbols and scale indices that `wai.image.jpegai` entropy-decodes (§13.6.1). The payload text four of these depend on, `entropy.json`, the `int_synth` payload and the `WIM1` payload with its `WMB1` bundle, is in [integer-payloads](integer-payloads.md), and §15.2 says where. The integer renderer of `wai.world.model` and the other integer networks of `wai.image.jpegai` are informative (§13.5, §13.6.2).
> **What a reader can build from this text.** With the payload extensions it cites, an implementation of this text decodes `wai.video.int_motion`, `wai.audio.int_codec`, `wai.splat.int_codec`, `wai.neural.int_synth`, `wai.neural.int_hyper`, `wai.video.int_hyper`, `wai.neural.int_mlicpp` and `wai.{image,video,audio,splat}.int_refine`, and reproduces the integers of `wai.image.jpegai`'s integer stages: the symbols and scale indices it entropy-decodes, and what its integer networks give for their inputs (§13.6). The payload text four of these depend on, `entropy.json`, the `int_synth` payload and the `WIM1` payload with its `WMB1` bundle, is in [integer-payloads](integer-payloads.md), and §15.2 says where. The integer renderer of `wai.world.model` is informative (§13.5).
>
> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`, `wai.splat.int_codec`, `wai.audio.int_codec`, `wai.neural.int_hyper`, `wai.video.int_hyper` and `wai.neural.int_mlicpp`, and the scale path and entropy decodes of `wai.image.jpegai` (§13.6.1), on both of its reference decode paths; and, through their stage operators alone with no graph (§13.5), `wai.video.int_motion`, the P-frames of `wai.video.int_hyper`, and `wai.{image,video,audio,splat}.int_refine`. The integer renderer of `wai.world.model` renders through it too, by `world/1` (§13.5), and stays informative. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ.
> Reference implementation: the `wai-ops` crate (`wai/wai-ops`): parse, lowering, interpreter, stage operators and bounds. Tool: `wai-verify bounds`. Conformance: `wai/ops-conformance/`. A capability's reference decoder decodes through this operator set from the revision that registers the capability's `ops` member (§11); in this revision, `wai.neural.int_synth`, `wai.splat.int_codec`, `wai.audio.int_codec`, `wai.neural.int_hyper`, `wai.video.int_hyper` and `wai.neural.int_mlicpp`, and the integer stages of `wai.image.jpegai` (§13.6): its scale path and entropy decodes on both of its reference decode paths, and its integer networks on its exact decoder; and, through their stage operators alone with no graph (§13.5), `wai.video.int_motion`, the P-frames of `wai.video.int_hyper`, and `wai.{image,video,audio,splat}.int_refine`. The integer renderer of `wai.world.model` renders through it too, by `world/1` (§13.5), and stays informative. Until then, a capability's reference implementation is the decoder its payload extension names, and §15.3 lists where this text and that decoder differ.
>
> The keywords **MUST**, **MUST NOT**, **SHOULD** and **MAY** are as defined in RFC 2119 and RFC 8174.
@ -198,8 +198,10 @@ An implementation forms each of these in a width that holds its certified interv
- no multiplier, one shift per channel, `pre` bias, `floor`: the three convolutions of the `wai.image.jpegai` scale path;
- no multiplier, shift 0, `pre` bias: the per-channel gain of that scale path;
- one multiplier per channel, one shift, `post` bias, `floor`: the convolutions of `WHM1` and `WAM1`;
- one multiplier, one shift, `post` bias, `floor`: the convolutions of the `wai.image.jpegai` PTQ networks;
- one multiplier, one shift, `post` bias, `floor`: the convolutions of the `wai.image.jpegai` integer networks;
- no multiplier, shift 16, no bias, `floor`: their Q16 products;
- no multiplier, shift 0, a `post` bias of `2^16`: their gates' `f + 1`;
- one multiplier `2^16`, shift 0, no bias: their hyper decoders' input;
- one multiplier, one shift, no bias, `half_up`: the RGB8 output step (multiplier 255) and the PCM output step (multiplier 32767);
- no multiplier, shift 16, no bias, `half_up`: `gram` and `apply` in channel attention;
- one multiplier per channel, shift 16, no bias, `half_up`: the attention temperature;
@ -1143,6 +1145,7 @@ The graph's identity is the SHA-256 of its bytes.
- `subset` is `full` or `npu`. A capability registered `npu` refuses a set that is not in the subset (§9.2).
- `lowering` names the §13 lowering, or `whg1` when the set carries graphs.
- `stage` fixes the modes of the stage operators the driver uses (§5E).
- `role_profiles`, OPTIONAL, names roles of the lowering (§13.0) whose graphs every conforming set holds to a stronger profile than `profile`, each with that profile. A graph meets a profile when its inputs' and scalars' declared ranges meet the profile's `tensor` width and every node of the graph fits the profile (§8.5). A set with a named graph that does not is refused (`model_mismatch`). In this revision `wai.image.jpegai` declares `profile` `acc64` and `role_profiles` `{ "hs_y": "acc32", "hs_uv": "acc32" }`: its scale path runs on a 32-bit fabric as on a 64-bit one, and its other integer networks meet `acc64` (§13.6.2).
- A capability whose decode runs no graph, only stage operators (§13.5), has no parameter set, so its member carries `set` and `stage` and no `profile`, `subset` or `lowering`. In this revision, `wai.video.int_motion` and `wai.{image,video,audio,splat}.int_refine`:
```json
@ -1193,7 +1196,7 @@ A lowering writes exactly one byte sequence for a set. These rules fix it, and
|---|---|---|---|---|
| `int_synth.json` | 4 | — | 4 | — |
| `WHM1`, `WAM1` | 2 | 8 | 8 | 4 |
| a `wai.image.jpegai` bundle (JSON) | smallest | — | smallest | — |
| a `wai.image.jpegai` bundle (JSON) | smallest | smallest | smallest | smallest |
**Roles and order.** A lowered set's graphs, by role, in this order:
@ -1205,7 +1208,7 @@ A lowering writes exactly one byte sequence for a set. These rules fix it, and
| `wai.splat.int_codec` | `whm1/1` | `g_s` | the RGB8 step |
| `wai.neural.int_mlicpp` | `wmb1/1` | `h_s`, `entropy`, `g_s` | in `g_s`, the RGB8 step |
| `wai.world.model` (informative) | `world/1` | `T`, `g_s` | in `g_s`, the RGB8 step |
| `wai.image.jpegai` | `jpegai-int/1` | `hs_y`, `hs_uv` | the per-channel gain |
| `wai.image.jpegai` | `jpegai-int/1` | `hs_y`, `hs_uv`, `rec_y`, `rec_uv`, `syn_y`, `syn_uv` | in `hs_*`, the per-channel gain; in `syn_*`, the denormalisation (§13.6.2) |
The **RGB8 step** is two operators after the network's last: `clamp(0, 2^afrac)`, then `rescale` with one multiplier 255 (`m_bytes` 4), one shift `afrac`, no bias and `half_up`. Everything else a driver does with an output, such as splitting `h_s`'s output into `σ` and `μ`, or laying values out as bytes, is a driver step and not an operator.
@ -1217,6 +1220,8 @@ The **RGB8 step** is two operators after the network's last: `clamp(0, 2^afrac)`
| a `whm1/1` or `wam1/1` graph | `x_0`, the network's input, with the policy and range of §13.2 | 0, 0 | `wam1/1` only: scalar 0, `t_out`, `[1, 2^26]` |
| `T` (`world/1`) | `z`, `check [−2^63, 2^63 − 1]` | 0, 0 | scalars `0 … A − 1`, the plane values, each `[−2^63, 2^63 − 1]` |
| `hs_y`, `hs_uv` | `z [Cz, 4, 4]`, `check [−31, 32]` | 4, 4 | none |
| `rec_y`, `rec_uv` | `z [C, 4, 4]`, `check [−31, 32]`; input 1, the residual `[C, 16, 16]`, `check [−2^31, 2^31 − 1]` | 4, 4; 16, 16 | none |
| `syn_y`, `syn_uv` | the luma latent `[160, 16, 16]`; for `syn_uv`, input 1, the chroma latent `[96, 16, 16]`; each `check [−2^31, 2^31 − 1]` | 16, 16 | none |
Sizes of 0 mean that the input's size is set per decode (§10.2). A driver plans these graphs on the payload's sizes, at the point its payload extension plans its networks (§4).
@ -1332,7 +1337,7 @@ Each capability's driver is defined by its payload extension. This table names t
### 13.6 The integer stages of `wai.image.jpegai` (`jpegai-int/1`)
The capability is `NeuralEntropyExact` (§12). §13.6.1 is its scale path and its two entropy decodes: normative, and complete from this text, the codestream syntax of ISO/IEC 6048-1, and the bundle members it names. §13.6.2 maps its other integer networks, and is informative.
The capability is `NeuralEntropyExact` (§12). §13.6.1 is its scale path and its two entropy decodes, and §13.6.2 its other integer networks. Both are normative, and complete from this text, the codestream syntax of ISO/IEC 6048-1, and the bundle members they name.
#### 13.6.1 The scale path and the entropy decodes
@ -1395,27 +1400,101 @@ The output is `v = I_s + g[c]`, `[Cz, 16, 16]`.
**What this subsection fixes.** `z`, `v`, the mask, `likely`, `σ_idx` and `y_q` are integers on every fabric. They are what the capability's `entropy-consistency` claim covers (§12).
#### 13.6.2 The other integer networks (informative)
#### 13.6.2 The hyper decoder, the context stages and the synthesis
The hyper decoder, the four context stages and the synthesis are integer as well. Their topology is not written out here, so this subsection maps each integer step to `wai-ops/1` and does no more.
The capability's other integer networks are each plane's hyper decoder, the luma plane's four context stages and each plane's synthesis. Their weights are integers and their values are Q16: an integer `q` stands for `q / 2^16`. `jpegai-int/1` lowers them to four graphs:
- `rec_y` and `rec_uv`: a plane's hyper decoder, and for luma its four context stages, from the plane's hyper-latent `z` and its dequantised residual to its reconstructed latent `ŷ`;
- `syn_y` and `syn_uv`: the synthesis, from the planes' latents to the planes.
| integer step | `wai-ops/1` |
**The members.** The graphs read the bundle's `weights/hyp_y.json` and `weights/hyp_uv.json` (the hyper decoders), `weights/mcm0_y.json` to `weights/mcm3_y.json` (the context stages 0 to 3), and `weights/synth_y.json` and `weights/synth_uv.json` (the synthesis). Each is a JSON object. A member whose value is an object with a `kind` of `"conv"` or `"q16"` is read, and every other member is ignored.
- Kind `"conv"`: `dims`, an array of non-negative integers; `w`, an array of as many integers as the product of `dims`, each fitting i64, row-major over `dims`; `mult`, an integer that fits i64; and `shift`, an integer in `0..=127`.
- Kind `"q16"`: `v`, an array of integers that fit i64. Its other members, `dims` among them, are not read.
A file in which any object, at any depth, has two members of one name, or in which a member of either kind lacks one of these or holds another value, is refused (`model_mismatch`). A member of neither kind is ignored, a bias among them: a module whose `m.bias` is of neither kind has no bias.
A **module** `m` names its members. Its weights are `m.weight`, of kind `conv`. Its bias is `m.bias`, of kind `q16`, and a module MAY have none. A layer normalisation's `m.weight` and `m.bias` are both of kind `q16`. A member that the graphs below read, and that the file lacks or holds as the other kind, is refused (`model_mismatch`). The member names are the bundle's, and mean nothing here.
**Blocks.** Each block reads a tensor `x` of `c` channels and appends operators, in the order given. `A32` is the policy `check [−2^31, 2^31 − 1]`.
| block | operators |
|---|---|
| the hyper decoder's input | `z·2^16`: `rescale` with one multiplier `2^16` and shift 0 |
| a convolution (grouped or depthwise) | `conv`; `RQ`: one multiplier (the convolution's own `mult`), one shift (its own `shift`, except as below), a `post` bias (its Q16 bias, or 0), `floor` |
| a transposed convolution | `conv_transpose`, with the same `RQ` |
| the shift inside a TAB | the attention half's three convolutions (`prep_data.conv1`, `prep_data.conv2`, `attn.project_out`) all use `prep_data.conv1`'s shift; the feed-forward half's three (`ffn.project_in`, `ffn.dwconv`, `ffn.project_out`) all use `ffn.project_in`'s. Each keeps its own multiplier |
| ReLU6, clip | `clamp(0, 6·2^16)`, `clamp` |
| a residual or mean addition | `add` |
| a Q16 product | `mul`; `RQ`: no multiplier, shift 16, `floor` |
| a gate `x·(f + 1)` in Q16 | `rescale` (a `post` bias of `2^16`), then `mul` |
| sigmoid, ELU, softmax | `sigmoid`, `elu(2^16)`, `softmax` |
| channel layer normalisation | `layernorm_c(ε = 42949)` |
| channel attention | `l2norm(42949)` on `q` and on `k`; `gram(heads, 16, half_up)`; `rescale` with one multiplier per channel (the head's temperature), shift 16 and `half_up`; `softmax`; `apply(heads, 16, half_up)` |
| channel slices; concatenation; the 2×2 sub-lattices, in the order `(0,0), (1,1), (0,1), (1,0)`; depth-to-space | `split`; `concat`; `subpos_split` and `subpos_merge`; `pixel_shuffle` |
| denormalisation | `rescale` with one multiplier 128, shift 0 and a `post` bias of `127·2^16 + 2^15`, with the policy `clamp [0, 255·2^16]`; for a byte, then `rescale` with shift 16 and `half_up` |
| `C(m; k, s, p, g)` | `conv` from `m.weight`, whose `dims` MUST be `[cout, c/g, k, k]` with `c` and `cout` multiples of `g`: `kh = kw = k`, `sh = sw = s`, `ph = pw = p`, `groups = g`. `g = dw` means `groups = c`, and then `cout = c`. Its `RQ`: `floor`; one multiplier, `mult`; one shift, `shift` (in a `TAB`, as below); a `post` bias of `cout` values from `m.bias`, or none when the file has no `m.bias`. The output is `A32`. |
| `T(m; k, op)` | `conv_transpose` from `m.weight`, whose `dims` MUST be `[c, cout, k, k]`: `kh = kw = k`, stride 2, padding 1, output padding `op`, with `C`'s `RQ`. The output is `A32`. |
| `relu6` | `clamp(0, 6·2^16)` |
| `R(m)` | `C(m.conv1; 3, 1, 1, 1)`, `relu`, `C(m.conv2; 3, 1, 1, 1)`, then `add(x, ·)`, `A32` |
| `G(m)` | `relu6`, `C(m.conv; 3, 1, 1, 4)`, `C(m.conv2; 1, 1, 0, 1)`, then `rescale` with no multiplier, shift 0, `floor` and a `post` bias of `2^16` on every channel, `A32`, then `mul(x, ·)`, `A32`: the gate `x·(f(x) + 1)` |
| `CAB` | on `x`: `R(CAB.residual_trunk.0)`, `R(CAB.residual_trunk.1)`, giving `t`; on `x` again: `C(CAB.subscale; 3, 2, 1, 1)`, `R(CAB.residual_mask2.0)`, `R(CAB.residual_mask3.0)`, `T(CAB.upscale; 3, 1)`, `sigmoid`, giving `s`; then `mul(t, s)`, `A32`, and `add(x, ·)`, `A32` |
| `D` | the denormalisation: `rescale` with one multiplier 128, shift 0, `floor` and a `post` bias of `127·2^16 + 2^15` on every channel, with the policy `clamp [0, 255·2^16]` |
Every other convolution uses its own shift. Today's bundle happens to give every convolution shift 31; the rule above is the decoder's for any bundle. These graphs' inputs come from floating-point stages, so once their topology is written, they have the policy `check`, with their admissible ranges for `acc64` (§8.4).
`mul` here is a Q16 product: no multiplier, shift 16, no bias, `floor`. The two inputs of each `mul` and each `add` have one channel count. A constant this lowering writes takes the smallest allowed width that holds it (§13.0).
**`TAB(m)`** reads `x` of `c` channels, `c` a multiple of 4, and `d = c/4`. Its attention half's convolutions take the shift of `m.prep_data.conv1.weight`, and its feed-forward half's that of `m.ffn.project_in.weight`. Each keeps its own multiplier. In order:
1. `layernorm_c` with `ε = 42949`, `γ` the `c` values of `m.prep_data.norm1.weight` and `β` those of `m.prep_data.norm1.bias`, `A32`;
2. `C(m.prep_data.conv1; 1, 1, 0, 1)`, whose `cout` MUST be `3c`, then `C(m.prep_data.conv2; 3, 1, 1, dw)`;
3. `split` at `0`, `c` and `2c`, each of `c` channels: `q`, `k` and `v`;
4. `l2norm` with `ε = 42949` on `q`, then on `k`;
5. `gram` of the two, with `groups` 4, shift 16 and `half_up`, `A32`;
6. `rescale` with one multiplier per channel, channel `i`'s the value at `⌊i/d⌋` of the 4 values of `m.attn.temperature`, shift 16, `half_up` and no bias, `A32`;
7. `softmax`, then `apply` of it to `v`, with `groups` 4, shift 16 and `half_up`, `A32`;
8. `C(m.attn.project_out; 1, 1, 0, 1)`, then `add(x, ·)`, `A32`, giving `x′`;
9. `layernorm_c` with `ε = 42949` from `m.ffn.norm1` on `x′`, `A32`;
10. `C(m.ffn.project_in; 1, 1, 0, 1)`, whose `cout` `2h` MUST be even, then `C(m.ffn.dwconv; 3, 1, 1, dw)`;
11. `split` at `0` and `h`, each of `h` channels: `a` and `b`;
12. `elu` with `α = 2^16` on `a`, then `mul(·, b)`, `A32`;
13. `C(m.ffn.project_out; 1, 1, 0, 1)`, then `add(x′, ·)`, `A32`.
Every other convolution takes its own shift. Today's bundle happens to give every convolution shift 31; the rule is the decoder's for any bundle.
**The graphs.** Each graph's inputs are listed in input-table order, each with the sizes it declares. The sub-lattice order is `(0,0), (1,1), (0,1), (1,0)`.
- **`rec_y`** (`C = 160`, from `hyp_y.json` and `mcm0_y.json` to `mcm3_y.json`) and **`rec_uv`** (`C = 96`, from `hyp_uv.json`). Input 0 is `z [C, 4, 4]`, `check [−31, 32]`, sizes 4 × 4: the hyper-latent of §13.6.1 step 1. Input 1 is the dequantised residual `r [C, 16, 16]`, `A32`, sizes 16 × 16. The hyper decoder comes first:
1. `rescale` of `z` with one multiplier `2^16`, shift 0, no bias and `floor`, its output `derived`;
2. `C(conv1; 1, 1, 0, 1)`, `T(conv2; 4, 0)`, `relu6`, `C(conv3; 3, 1, 1, 1)`, `relu6`, `C(conv4; 3, 1, 1, 1)`, whose `cout` MUST be `4C`. Its output is `ψ`.
Then, for `rec_uv`: `split(ψ, iC, C)` for `i = 0 … 3`; `subpos_merge` of the four in the sub-lattice order; and `add(r, ·)`, `A32`, which is the graph's output `ŷ`.
For `rec_y`: `subpos_split(r)` in the sub-lattice order, giving `r_0 … r_3`. Then, for each context stage `i = 0 … 3`, from `mcm{i}_y.json`:
1. `split(ψ, iC, C)`, giving `p_i`;
2. for `i = 0`, `f = p_i`. For `i ≥ 1`, with the context `y_0` for `i = 1` and `concat(y_0, …, y_{i−1})` otherwise: `C(conv.0; 1, 1, 0, 1)`, then `C(conv.1; 3, 1, 1, 5)`, then `concat(·, p_i)`, giving `f`;
3. on `f`: `C(fusion_pred_net.conv1; 1, 1, 0, 1)`, `relu`, `C(fusion_pred_net.conv2; 1, 1, 0, 1)`, `relu`, `C(fusion_pred_net.conv3; 1, 1, 0, 1)`, whose `cout` MUST be `C`, giving `μ_i`;
4. `add(r_i, μ_i)`, `A32`, giving `y_i`.
Last, `subpos_merge(y_0, y_1, y_2, y_3)` in the sub-lattice order, which is the graph's output `ŷ`.
- **`syn_y`**, from `synth_y.json`. Input 0 is the luma latent `[160, 16, 16]`, `A32`, sizes 16 × 16. Then: `R(first_stage.0.0)`, `T(first_stage.1; 3, 1)`, `G(first_stage.2)`, `T(conv2_t; 3, 1)`, `CAB`, `G(iact2)`, `C(conv3_t; 1, 1, 0, 1)`, `pixel_shuffle(2)`, `C(TAM.ds_conv; 3, 2, 1, 1)`, `TAB(TAM.TABs.0)`, `TAB(TAM.TABs.1)`, `T(TAM.us_conv; 3, 1)`, `G(iact3)`, `T(conv4_t; 3, 1)`, whose `cout` MUST be 1, and `D`. The output is the luma plane `[1, 256, 256]`.
- **`syn_uv`**, from `synth_uv.json`. Input 0 is the luma latent `[160, 16, 16]` and input 1 the chroma latent `[96, 16, 16]`, each `A32`, sizes 16 × 16. Then: `concat` of inputs 0 and 1; `C(first_stage.conv1; 3, 1, 1, 1)`; `concat` of that and input 1; `T(conv2_t; 3, 1)`, `CAB`, `G(iact2)`, `C(conv3_t; 1, 1, 0, 1)`, `pixel_shuffle(2)`, `TAB(TAM.TABs.0)`, `TAB(TAM.TABs.1)`, `G(iact3)`, `T(conv4_t; 3, 1)`, whose `cout` MUST be 8, `pixel_shuffle(2)`, and `D`. The output is the chroma planes `[2, 256, 256]`, U then V.
The input of each `pixel_shuffle(2)` has a multiple of 4 channels. Every tensor that no block above gives `A32` or `clamp` is `derived`: the outputs of `clamp`, `relu`, `sigmoid`, `softmax`, `l2norm`, `elu` and the movement operators, and the hyper decoder's scaled input. Each of those lies within the hull of its inputs' intervals and `[−2^16, 2^16]`, and the scaled input within `[−31·2^16, 32·2^16]`.
**Range checks.** These graphs check every activation that an operator can grow, at `A32`, and so every sum they form is bounded by the weights alone. Their inputs come from floating-point stages, so these checks refuse what the direct decoder computed on:
- a dequantised residual or a latent whose Q16 value leaves i32 (input 1 of `rec_*`, the inputs of `syn_*`);
- an activation that leaves i32: the output of a convolution, a transposed convolution, an addition, a Q16 product, a layer normalisation, an attention contraction, the attention scale or a gate's `f + 1`.
Each is `field_range` (§4).
**Load checks.** In addition to §13.6.1's, a bundle is refused (`model_mismatch`) when a file breaks the member rules above; when a block's `dims` or value counts are not those it names, or a channel count is not a multiple it names; when a graph does not plan (§10.5 step 7); or when the certificate does not meet `acc64`.
**Driver.** For each plane, after §13.6.1's steps:
1. The floating-point stages dequantise `y_q` into the residual. The driver converts each residual value `x` to the integer nearest `x·2^16`, a tie rounding away from zero, and runs `rec_*` on `z` and that residual. It gives `ŷ`.
2. The floating-point stages read `ŷ` as `ŷ / 2^16`, refine it into the plane's latent, and the driver converts the latent to Q16 as in step 1.
Then `syn_y` runs on the luma latent, and `syn_uv` on both latents. The post-filters read each output value `v` as `v / 2^16`.
**Refusals.** A bundle that fails a load check: `model_mismatch`. A decode whose input or activation leaves its range: `field_range`. [determinism-tiers](determinism-tiers.md) says how the capability reports them. A refused exact decode MUST NOT fall back silently.
**Profiles.** `layernorm_c`, `l2norm`, `softmax`, `sigmoid` and `elu` form values above i32 whatever their inputs (§7.3), so these four graphs meet `acc64` at best. The capability's `ops` member declares `acc64` for the set, and `acc32` for the scale path's graphs `hs_y` and `hs_uv` (§11).
**What this subsection fixes.** Given the same `z` and residual, `ŷ` is the same integers on every fabric whose implementation conforms to `acc64`. Given the same latents, so are the synthesis planes. The inputs come from floating-point stages, so the capability's claim stays `entropy-consistency` (§12): these graphs MAY run on an implementation that does not conform, and their output is then outside the claim.
**Accumulators, measured (informative).** No sum of these graphs needs more than 64 bits. On the reference codestream of `jpegai/v1` (`tests/jpegai_nets_ops.rs`), every value each operator forms lies within its certified interval, and the largest magnitudes reached are:
| graph | `conv.acc` | `convT.acc` | `mul.p` | `gram.acc` | `apply.acc` | an input | an activation |
|---|---|---|---|---|---|---|---|
| `rec_y` | 36 (53) | 33 (48) | — | — | — | 19 | 21 |
| `rec_uv` | 33 (42) | 32 (46) | — | — | — | 17 | 21 |
| `syn_y` | 36 (55) | 35 (54) | 35 (60) | 32 (46) | 30 (49) | 21 | 24 |
| `syn_uv` | 37 (55) | 36 (52) | 37 (58) | 32 (46) | 33 (49) | 21 | 24 |
Each cell is `bits` (§2) of the largest magnitude reached under that key in that graph, with the certificate's `bits` for the key in brackets. The activations include the denormalised planes, which lie in `[0, 255·2^16]`. The certified bounds come from `A32` and the bundle's weights. No measured sum and no certified one reaches `2^63`, so `acc64` holds for every sum of these graphs, and no wider accumulator is needed. The `product` intermediates, which `acc64` forms in 128 bits (§7.1), reach 53 bits (`rq.p`).
## 14. Conformance
@ -1539,7 +1618,7 @@ These land with each capability's lowering, capability by capability:
- **SPEC §3.1:** the OPTIONAL pin member `ops_profile` (§11).
- **SPEC §5:** the rows of `wai.neural.int_synth` and `wai.neural.int_mlicpp` cite integer-payloads §13 and §14, and §13.1 and §13.4.
- **SPEC §7:** a `NeuralIntegerExact` capability names its `ops` member.
- **determinism-tiers:** §5 cites §12; the execution fields carry the profile; `wai.image.jpegai`'s exact decoder reports `model_mismatch` and `stream_exhausted` (§13.6.1) and does not fall back silently.
- **determinism-tiers:** §5 cites §12; the execution fields carry the profile; `wai.image.jpegai`'s exact decoder reports `model_mismatch` and `stream_exhausted` (§13.6.1) and `field_range` (§13.6.2), and does not fall back silently.
- **staged-delivery A.4:** cites E9, E11 and E12.
- **`registry/capabilities.json`:** the `ops` members.
@ -1566,6 +1645,7 @@ No committed output changes, and no committed refusal fixture changes its code.
4. **`wai.image.jpegai`.** These take effect for its exact decoder and for the reference sink's float decode path, which decode the scale path and the entropy through the same lowering.
- Its exact decoder refuses a malformed latent or hyper-latent stream (`stream_exhausted`) where it panicked.
- It refuses a bundle that fails §13.6.1's load checks (`model_mismatch`), including a bound above 255 or a shift outside `0..=127`, which its loader wrapped; and a bundle whose scale graphs do not meet `acc32`, the profile its `ops` member declares (§11): a sum past `i32`, or a shift past 63, which it computed in `i64` or shifted as its build did.
- The codestream's scope is checked first, then the bundle, then the streams. A refusal of the exact decode is its result: a sink that holds the bundle does not then decode the capability through the float path ([determinism-tiers](determinism-tiers.md) §5).
- The codestream's scope is checked first, then the bundle (its scale path, then its other networks), then the streams, then the networks' ranges as each runs. A refusal of the exact decode is its result: a sink that holds the bundle does not then decode the capability through the float path ([determinism-tiers](determinism-tiers.md) §5).
- Its residual-variance bands start at `likely` 1004, 1353 and 1938 (§13.6.1 step 3), where the reference decoder's own tables change value. An earlier revision of this text and of the exact decoder started them one lower, so a block whose `likely` was exactly 1003, 1352 or 1937 took the next band's addition and dequantisation multiplier, where the reference decoder and the reference sink's float decode path, which reads the reference decoder's tables, did not. No committed codestream has such a block.
- Its exact decoder runs its hyper decoders, context stages and synthesis through `rec_y`, `rec_uv`, `syn_y` and `syn_uv` (§13.6.2). It refuses a bundle whose weight files break §13.6.2's member rules, whose blocks do not have the shapes the topology names, or whose certificate does not meet `acc64` (`model_mismatch`), where its loaders panicked on a mistyped member or it computed with what the files held. It refuses a decode in which a dequantised residual, a latent or an activation of these networks leaves i32 as a Q16 value (`field_range`), where it computed on in 64 or 128 bits, or wrapped or panicked past them. No committed bundle or codestream reaches either: on the reference codestream every input and activation lies within i32 (§13.6.2, *Accumulators, measured*).
5. **`WIV1` and `WIR1`.** None. Their decodes through the stage operators (§13.5) give the output and the refusal their direct decoders gave, on every input: the operators are the arithmetic integer-payloads §3 and staged-delivery Appendix A already defined, and the refusals keep their codes and their order.

2
wai/ops-conformance/certificates/external.json
View file

@ -1 +1 @@
{"format":"wai-ops-external/1","sets":[{"capability":"wai.neural.int_hyper","lowering":"whm1/1","members":{"entropy.json":"cond_entropy.json","g_s.whm":"int_synth_sota_model.bin","h_s.whm":"int_hs_model.bin"},"name":"sota","note":"regenerated from pinned inputs by the job that certifies it; not committed","profile":"acc64","set_sha256":"f16fd09cd79b831d7d821e030c05e99e9fb5ea216574e699944ea4e55ac0a404","workflow":".github/workflows/wai-sota-vectors.yml"},{"capability":"wai.image.jpegai","lowering":"jpegai-int/1","members":{"norm.json":"norm.json","weights/hs_uv.json":"hs_uv.json","weights/hs_y.json":"hs_y.json","weights/hyp_uv.json":"hyp_uv.json","weights/hyp_y.json":"hyp_y.json","weights/icci.json":"icci.json","weights/mcm0_y.json":"mcm0_y.json","weights/mcm1_y.json":"mcm1_y.json","weights/mcm2_y.json":"mcm2_y.json","weights/mcm3_y.json":"mcm3_y.json","weights/synth_uv.json":"synth_uv.json","weights/synth_y.json":"synth_y.json"},"name":"jpegai/v1","note":"the models root's jpegai/v1, the files tools/jpegai-e2e-v1.sha256 pins; not committed","profile":"acc32","set_sha256":"3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37","workflow":".github/workflows/wai-jpegai-e2e.yml"}]}
{"format":"wai-ops-external/1","sets":[{"capability":"wai.neural.int_hyper","lowering":"whm1/1","members":{"entropy.json":"cond_entropy.json","g_s.whm":"int_synth_sota_model.bin","h_s.whm":"int_hs_model.bin"},"name":"sota","note":"regenerated from pinned inputs by the job that certifies it; not committed","profile":"acc64","set_sha256":"f16fd09cd79b831d7d821e030c05e99e9fb5ea216574e699944ea4e55ac0a404","workflow":".github/workflows/wai-sota-vectors.yml"},{"capability":"wai.image.jpegai","lowering":"jpegai-int/1","members":{"norm.json":"norm.json","weights/hs_uv.json":"hs_uv.json","weights/hs_y.json":"hs_y.json","weights/hyp_uv.json":"hyp_uv.json","weights/hyp_y.json":"hyp_y.json","weights/icci.json":"icci.json","weights/mcm0_y.json":"mcm0_y.json","weights/mcm1_y.json":"mcm1_y.json","weights/mcm2_y.json":"mcm2_y.json","weights/mcm3_y.json":"mcm3_y.json","weights/synth_uv.json":"synth_uv.json","weights/synth_y.json":"synth_y.json"},"name":"jpegai/v1","note":"the models root's jpegai/v1, the files tools/jpegai-e2e-v1.sha256 pins; not committed","profile":"acc64","role_profiles":{"hs_uv":"acc32","hs_y":"acc32"},"set_sha256":"3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37","workflow":".github/workflows/wai-jpegai-e2e.yml"}]}

2
wai/ops-conformance/certificates/index.json
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9.json Normal file
View file

File diff suppressed because one or more lines are too long

0
wai/ops-conformance/certificates/wai.image.jpegai/60020d0657c4561db336a42f2d359b254a3f12348e8122aeb7bbb91d63b180ad/hs_uv.whg1 → wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/hs_uv.whg1
View file

0
wai/ops-conformance/certificates/wai.image.jpegai/60020d0657c4561db336a42f2d359b254a3f12348e8122aeb7bbb91d63b180ad/hs_y.whg1 → wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/hs_y.whg1
View file

BIN
wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/rec_uv.whg1 Normal file
View file

Binary file not shown.

BIN
wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/rec_y.whg1 Normal file
View file

Binary file not shown.

BIN
wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/syn_uv.whg1 Normal file
View file

Binary file not shown.

BIN
wai/ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/syn_y.whg1 Normal file
View file

Binary file not shown.

2
wai/ops-conformance/certificates/wai.image.jpegai/3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37.json
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/certificates/wai.image.jpegai/60020d0657c4561db336a42f2d359b254a3f12348e8122aeb7bbb91d63b180ad.json
View file

File diff suppressed because one or more lines are too long

8
wai/ops-conformance/graphs/lower_jpegai.jsonl
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/hyp_uv.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/hyp_y.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/mcm0_y.json Normal file
View file

@ -0,0 +1 @@
{"fusion_pred_net.conv1.weight":{"kind":"conv","dims":[2,160,1,1],"mult":204,"shift":29,"w":[-3716,19116,-1860,-24856,32678,-11115,1556,9660,21290,-18707,27458,26250,-15008,-12356,24616,-29890,-31012,-21607,-2059,-26727,-2988,-18811,4887,-17302,-9819,29687,-15722,8405,-5965,-29422,-9181,20123,288,31495,789,-29567,-29616,20939,-9567,-10105,9613,-26804,5627,-26911,-19424,4293,-3830,-14516,-19614,24520,12482,-17256,-2032,-30370,24229,29069,-19504,31665,8313,3645,13036,13691,14002,-7584,18744,20004,23606,-26303,-15196,14226,-14391,-21791,18274,-2016,-8915,23346,-8133,15251,-22441,31842,27127,10304,28098,-25676,-16340,-347,12738,7723,29695,32467,17302,2816,-30378,5870,21905,-19552,22028,-14632,-3535,20980,21621,3425,16351,-26963,6213,25030,-17318,-17699,2040,10578,2034,15521,14460,29397,8083,-16237,-2665,-5384,10810,12835,24958,11102,1835,1792,19393,4956,-6843,-19397,-2103,-24626,24054,30084,29267,10775,-25393,14165,23356,-285,-10368,4852,-30268,-537,9412,-16726,-18249,-29751,4182,24545,-1900,-746,-27711,9361,8995,16590,-2519,29775,25611,15291,-19525,25955,-24160,26936,13542,32218,-21834,-10134,-10353,-1132,-12307,2563,-30620,-29420,-24412,-26086,24943,24125,-14757,-27504,1228,32617,-32732,28804,-15003,2639,4117,25423,14272,25860,28822,23984,2789,-12582,-2553,-9474,9358,-10641,-433,-28448,14580,-2811,-3183,-27430,24971,-1731,6186,-22358,-9283,-28800,9518,24518,-7048,24623,17481,-31446,-29347,-3357,-5616,32061,19477,-17761,-28014,18032,9405,32167,-18825,-27922,-2252,18536,16929,2722,19507,-10504,-3800,-27983,-6553,31275,23482,2970,7405,-30816,24887,25212,22459,-24965,23879,5283,8145,-2291,-32301,-7240,-23975,-2797,2346,8900,28776,24150,-32445,3992,17335,8155,-10173,19796,-22851,16947,21837,3182,-13147,21335,4823,31722,31652,-29243,293,29442,-14368,32274,-9938,12899,3291,-27402,-23778,-28138,-2211,2404,-12203,-6593,-14108,-9521,30060,-16393,-5691,24828,-2746,6970,-8713,-11369,-7237,-31708,-19389,10865,19938,-23787,14339,30552,26830,-10405,-4163,4258,-20809,19442,-21867,-14232,23825,25712,-17211,-29516,-4218,12690,-18599,32427]},"fusion_pred_net.conv2.weight":{"kind":"conv","dims":[2,2,1,1],"mult":16384,"shift":29,"w":[30611,15983,-22516,-1064]},"fusion_pred_net.conv3.weight":{"kind":"conv","dims":[160,2,1,1],"mult":16384,"shift":29,"w":[-11517,31463,-8565,19044,28589,9055,1812,17682,24197,15593,29157,28924,17526,-4427,-22507,-12660,-19308,24257,-3869,19378,-26961,-30432,17615,9863,-16831,21614,-12667,-32265,-21587,28586,-14229,25487,-28187,-25295,-28708,5740,-28550,19573,-27625,-12151,32466,-13872,6731,-20650,-29165,19466,-13785,-27911,9571,24282,2317,-11222,-3096,-30421,12603,-15952,-16571,13079,-22180,-3901,17885,23484,13464,12848,3793,-5712,3866,-6774,-12885,-23662,17921,-4974,-2638,23764,-19950,15871,-20574,-9144,10025,13508,-8615,17372,-4091,-1864,-12570,25658,-27639,29500,7937,27543,-23627,10806,17981,1582,19115,-16506,29243,13646,11993,-21447,10171,-251,23059,1336,-7829,-16404,4963,3454,-14777,18345,-17137,30046,-3197,-10881,12055,1983,-26960,-8598,32328,20657,-24864,6285,-29133,2023,-14469,8642,25298,-820,25201,-9837,16162,-26663,26322,32684,32284,9693,28381,-5960,22514,11750,6255,30320,-29937,7059,-22147,-27624,-27382,-18068,9102,18158,-15198,2489,-4024,25348,-23051,-26757,-8192,-17626,-10396,-9977,-6553,5500,24262,-25308,18169,14942,18627,-194,-31073,-1986,-2977,15865,1615,18558,-4063,-10174,-5588,-2855,16757,-29243,-28094,23426,26537,-2477,-25493,-23629,15799,-22618,-25469,-30980,-26444,-3314,-12884,3784,25946,29817,-8045,7323,4029,20048,-27741,-8809,32723,26182,-199,16300,-29868,4225,21548,5760,-16083,9958,-15908,3107,23102,-25778,2011,18659,-4026,2997,-19702,31345,8304,-2372,-1646,15729,-14480,12989,-3423,-250,15321,12590,-1907,-24798,-8683,24055,-26658,6147,-4836,22846,-3772,29035,9914,6550,7916,-8159,21710,5313,10018,-6948,16828,13830,-22242,-28320,-27973,-25443,9323,-24022,7818,28971,20897,-2190,16059,32508,-28211,2713,-14178,23152,-23326,-25550,13272,22770,-10303,-16053,15449,-1066,30199,17855,-12509,15465,20748,-28373,20183,7336,-5371,27941,-403,-31483,1558,-25559,17758,3380,-25123,-3348,-26824,27778,20570,9633,-30145,-15398,9050,27690,-25770,13154,-22980,14896,14888,-1742,-18629,16760,29514,-8700,-15825,-9710,26336,10268,-19122,-10251,31067,-2028]}}

1
wai/ops-conformance/sets/jpegai/weights/mcm1_y.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/mcm2_y.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/mcm3_y.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/synth_uv.json Normal file
View file

File diff suppressed because one or more lines are too long

1
wai/ops-conformance/sets/jpegai/weights/synth_y.json Normal file
View file

File diff suppressed because one or more lines are too long

8
wai/registry/capabilities.json
View file

@ -547,9 +547,13 @@
"registered": "2026-06-13",
"ops": {
"set": "wai-ops/1",
"profile": "acc32",
"profile": "acc64",
"subset": "full",
"lowering": "jpegai-int/1"
"lowering": "jpegai-int/1",
"role_profiles": {
"hs_y": "acc32",
"hs_uv": "acc32"
}
}
},
{

14
wai/registry/capabilities.schema.json
View file

@ -152,6 +152,9 @@
"lowering": [
"profile",
"subset"
],
"role_profiles": [
"profile"
]
},
"anyOf": [
@ -185,6 +188,17 @@
"lowering": {
"type": "string"
},
"role_profiles": {
"type": "object",
"description": "Roles of the lowering whose graphs every conforming set holds to a stronger profile than `profile`, each with that profile (extensions/integer-operators.md §11).",
"minProperties": 1,
"additionalProperties": {
"enum": [
"acc32",
"acc64"
]
}
},
"stage": {
"type": "object",
"additionalProperties": false,

456
wai/tools/wai_ops_check.py
View file

@ -23,7 +23,9 @@ extension's text (and, for the source forms it lowers, integer-payloads), it:
for every set of index.json and the two falsification sets of section 14.3,
each of which must fail the profile it claims. The sets of external.json are
certified too when their files are given (--external NAME=DIR); otherwise their
certificates are checked for what needs no set.
certificates are checked for what needs no set. Either way, an entry's verdict
must meet its `profile`, and each graph its `role_profiles` names must meet its
own profile (section 11).
python3 -I tools/wai_ops_check.py ops-conformance # check
python3 -I tools/wai_ops_check.py ops-conformance --write DIR # also write the files
@ -32,14 +34,25 @@ extension's text (and, for the source forms it lowers, integer-payloads), it:
Exit status 1 on any difference or failure.
python3 -I tools/wai_ops_check.py certify --capability CAP --lowering ID \\
[--c N] --member NAME=PATH ... [--profile P] [--out FILE] [--check FILE]
[--c N] --member NAME=PATH ... [--profile P] [--role-profile ROLE=P ...] \\
[--out FILE] [--check FILE]
python3 -I tools/wai_ops_check.py certify --whg1 ROLE=PATH ... [--profile P] ...
certifies one set, or WHG1 graphs on their own, writing the certificate to FILE or
to stdout. Exit status: 0 the verdict meets the profile (acc64 by default); 1 it
does not, and the certificate is still written; 2 a graph is malformed or refused
at load; 3 --check finds a difference, and the first differing path is printed with
both values; 4 the lowering refuses the set; 5 usage.
to stdout. --role-profile (repeatable) holds the graph of one role to a profile of
its own, as a capability's `role_profiles` does (section 11): that graph meets it
when its inputs' and scalars' declared ranges meet the profile's tensor width and
every node of it fits the profile (section 8.5). Exit status: 0 the verdict meets
the profile (acc64 by default) and each named graph meets its own; 1 one does not,
and the certificate is still written; 2 a graph is malformed or refused at load;
3 --check finds a difference, and the first differing path is printed with both
values; 4 the lowering refuses the set; 5 usage, or a --role-profile role that
names no graph.
The lowering jpegai-int/1 (section 13.6) reads eleven members, norm.json and
weights/{hs_y,hs_uv,hyp_y,hyp_uv,mcm0_y,mcm1_y,mcm2_y,mcm3_y,synth_y,synth_uv}.json,
and writes six roles: hs_y and hs_uv (13.6.1), rec_y, rec_uv, syn_y and syn_uv
(13.6.2). Any other member is part of the set's digest and is not read.
"""
import hashlib
import json
@ -1595,10 +1608,13 @@ def canonical(v):
return cw(v).encode("ascii")
def certificate(capability, lowering, set_sha256, roles):
def certificate(capability, lowering, set_sha256, roles, role_profiles=None):
"""The bounds certificate of section 8.5 for the graphs [(role, WHG1 bytes)], in
the order given: (its canonical bytes, its verdict). Each graph passes load steps
1-7 and the bounds pass, or the Refuse of the step that fails propagates."""
1-7 and the bounds pass, or the Refuse of the step that fails propagates. With
`role_profiles` ({role: profile}, section 11), it returns a third value, the
failure of section 8.5 of each named graph that does not meet its profile, judged
on that graph alone ({role: failure}; empty when every one meets it)."""
gs = [(name, b, load_graph(b)) for name, b in roles]
v = verdict_of([(name, g) for name, _, g in gs])
cert = {
@ -1609,7 +1625,17 @@ def certificate(capability, lowering, set_sha256, roles):
"set_sha256": set_sha256,
"verdict": v,
}
return canonical(cert), v
if role_profiles is None:
return canonical(cert), v
fails = {}
for name, p in role_profiles.items():
# A graph meets a profile when its inputs' and scalars' declared ranges meet
# the profile's tensor width and every node of it fits the profile: the
# verdict of section 8.5 over that graph alone.
rv = verdict_of([(n, g) for n, _, g in gs if n == name])
if not rv[p]:
fails[name] = rv[p + "_first_failure"]
return canonical(cert), v, fails
ABSENT = "(absent)"
@ -2118,6 +2144,12 @@ class WB:
self.ops.append((code, ins, [t], params))
return t
def op_n(self, code, ins, out_recs, params):
"""An operator of several outputs, which take the next indices in their order."""
ts = [self.tensor(*r) for r in out_recs]
self.ops.append((code, ins, ts, params))
return ts
def bytes(self, outputs):
b = bytearray(b"WHG1")
b += bytes([1, 1, 0, 0])
@ -2454,8 +2486,346 @@ def world_lower(tb, gb, c):
return [("T", wb.bytes([out])), ("g_s", net_graph(gn, True))]
def jpegai_lower(norm_b, hy_b, huv_b):
"""jpegai-int/1, section 13.6.1."""
JPEGAI_MEMBERS = (
"norm.json", "weights/hs_y.json", "weights/hs_uv.json", "weights/hyp_y.json", "weights/hyp_uv.json",
"weights/mcm0_y.json", "weights/mcm1_y.json", "weights/mcm2_y.json", "weights/mcm3_y.json",
"weights/synth_y.json", "weights/synth_uv.json",
)
def jpegai_lower(files):
"""jpegai-int/1 (section 13.6): the scale graphs of 13.6.1, then the graphs of
13.6.2, in section 13.0's role order. `files` maps member names to bytes; a member
the lowering reads that the set lacks refuses it. Other members are not read."""
if any(m not in files for m in JPEGAI_MEMBERS):
lref()
roles = jpegai_scale_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
W = {k: q16_file(files["weights/%s.json" % k]) for k in ("hyp_y", "hyp_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "synth_y", "synth_uv")}
roles.append(("rec_y", rec_graph(W["hyp_y"], [W["mcm%d_y" % i] for i in range(4)], 160)))
roles.append(("rec_uv", rec_graph(W["hyp_uv"], None, 96)))
roles.append(("syn_y", syn_y_graph(W["synth_y"])))
roles.append(("syn_uv", syn_uv_graph(W["synth_uv"])))
return roles
def q16_file(b):
"""One weights file of section 13.6.2, *The members*: {name: member} for each
member whose value is an object of kind "conv" or "q16"; every other member is
ignored. Refused: a file that is not a JSON object, two members of one object with
one name (anywhere in the file), and a member of either kind that lacks one of its
fields or holds another value."""
def pairs(ps):
if len({k for k, _ in ps}) != len(ps):
lref()
return dict(ps)
try:
root = json.loads(b, object_pairs_hook=pairs)
except ValueError:
lref()
if not isinstance(root, dict):
lref()
isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
out = {}
for name, m in root.items():
if not isinstance(m, dict) or m.get("kind") not in ("conv", "q16"):
continue
if m["kind"] == "conv":
dims, w, mult, shift = m.get("dims"), m.get("w"), m.get("mult"), m.get("shift")
if not isinstance(dims, list) or not all(isint(d) and d >= 0 for d in dims):
lref()
n = 1
for d in dims:
n *= d
if not isinstance(w, list) or len(w) != n or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in w):
lref()
if not isint(mult) or not I64_MIN <= mult <= I64_MAX or not isint(shift) or not 0 <= shift <= 127:
lref()
out[name] = {"kind": "conv", "dims": dims, "w": w, "mult": mult, "shift": shift}
else:
v = m.get("v")
if not isinstance(v, list) or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in v):
lref()
out[name] = {"kind": "q16", "v": v}
return out
A32 = (2, I32_MIN, I32_MAX)
Q16_RQ = rq_block(0, 0, [], 0, 0, [16], 1, [], 0) # mul: no multiplier, shift 16, no bias, floor
ATT_RQ = rq_block(1, 0, [], 0, 0, [16], 1, [], 0) # gram, apply: shift 16, half_up
SUBLATTICE = bytes([0, 3, 1, 2]) # (0,0), (1,1), (0,1), (1,0), each 2a + b
class Q16Net:
"""The blocks of section 13.6.2 over one graph under construction, reading the
modules of the weights file F (which a graph may switch, as rec_y does)."""
def __init__(self):
self.wb = WB()
self.F = {}
def c(self, x):
return self.wb.tensors[x][0]
def weight(self, m):
e = self.F.get(m + ".weight")
if e is None or e["kind"] != "conv":
lref()
return e
def q16(self, name):
e = self.F.get(name)
if e is None or e["kind"] != "q16":
lref()
return e["v"]
def bias(self, m, cout):
"""m.bias, cout values, or None when the file has none."""
e = self.F.get(m + ".bias")
if e is None:
return None
if e["kind"] != "q16" or len(e["v"]) != cout:
lref()
return e["v"]
def rq(self, e, shift, bias):
"""C's RQ: floor; one multiplier; one shift; a post bias, or none."""
mult = e["mult"]
if bias is None:
return rq_block(0, 0, [mult], 1, smallest([mult], (4, 8)), [shift], 1, [], 0)
return rq_block(0, 2, [mult], 1, smallest([mult], (4, 8)), [shift], 1, bias, smallest(bias, (4, 8)))
def wbytes(self, w):
if any(not I64_MIN <= v <= I64_MAX for v in w):
lref()
return smallest(w, (1, 2, 4, 8))
def conv(self, x, m, k, s, p, g, shift=None):
"""C(m; k, s, p, g); g = "dw" is groups = c, with cout = c."""
c = self.c(x)
e = self.weight(m)
dims = e["dims"]
dw = g == "dw"
if dw:
g = c
if len(dims) != 4 or c % g:
lref()
cout = dims[0]
if cout == 0 or cout % g or dims != [cout, c // g, k, k] or (dw and cout != c):
lref()
rq = self.rq(e, e["shift"] if shift is None else shift, self.bias(m, cout))
w = e["w"]
return self.wb.op(0x01, [x], (cout,) + A32, p_conv(cout, k, k, s, s, p, p, g, self.wbytes(w), rq, w))
def convt(self, x, m, k, op_):
"""T(m; k, op): stride 2, padding 1, output padding op, with C's RQ."""
c = self.c(x)
e = self.weight(m)
dims = e["dims"]
if len(dims) != 4 or dims[0] != c or dims[1] == 0 or dims[2:] != [k, k]:
lref()
cout = dims[1]
rq = self.rq(e, e["shift"], self.bias(m, cout))
w = e["w"]
return self.wb.op(0x11, [x], (cout,) + A32, p_convt(cout, k, k, 2, 2, 1, 1, op_, op_, self.wbytes(w), rq, w))
def same(self, *xs):
"""The inputs of each mul and each add have one channel count."""
if len({self.c(x) for x in xs}) != 1:
lref()
return self.c(xs[0])
def add(self, a, b):
return self.wb.op(0x05, [a, b], (self.same(a, b),) + A32, b"")
def mul(self, a, b):
return self.wb.op(0x10, [a, b], (self.same(a, b),) + A32, Q16_RQ)
def unary(self, code, x, params=b""):
return self.wb.op(code, [x], (self.c(x),), params)
def relu6(self, x):
return self.unary(0x0B, x, p_clamp(0, 6 << 16))
def split(self, x, start, count):
if start + count > self.c(x):
lref()
return self.wb.op(0x08, [x], (count,), struct.pack("<II", start, count))
def concat(self, xs):
return self.wb.op(0x06, list(xs), (sum(self.c(x) for x in xs),), b"")
def shuffle2(self, x):
if self.c(x) % 4:
lref()
return self.wb.op(0x09, [x], (self.c(x) // 4,), struct.pack("<I", 2))
def merge(self, xs):
return self.wb.op(0x16, list(xs), (self.same(*xs),), SUBLATTICE)
def R(self, x, m):
t = self.conv(x, m + ".conv1", 3, 1, 1, 1)
t = self.unary(0x0C, t)
t = self.conv(t, m + ".conv2", 3, 1, 1, 1)
return self.add(x, t)
def G(self, x, m):
"""The gate x·(f(x) + 1)."""
t = self.relu6(x)
t = self.conv(t, m + ".conv", 3, 1, 1, 4)
t = self.conv(t, m + ".conv2", 1, 1, 0, 1)
c = self.c(t)
t = self.wb.op(0x02, [t], (c,) + A32, rq_block(0, 2, [], 0, 0, [0], 1, [1 << 16] * c, 4))
return self.mul(x, t)
def CAB(self, x):
t = self.R(x, "CAB.residual_trunk.0")
t = self.R(t, "CAB.residual_trunk.1")
s = self.conv(x, "CAB.subscale", 3, 2, 1, 1)
s = self.R(s, "CAB.residual_mask2.0")
s = self.R(s, "CAB.residual_mask3.0")
s = self.convt(s, "CAB.upscale", 3, 1)
s = self.unary(0x1B, s)
return self.add(x, self.mul(t, s))
def D(self, x):
"""The denormalisation, clamp [0, 255·2^16]."""
c = self.c(x)
rq = rq_block(0, 2, [128], 1, 4, [0], 1, [127 * (1 << 16) + (1 << 15)] * c, 4)
return self.wb.op(0x02, [x], (c, 1, 0, 255 << 16), rq)
def layernorm(self, x, m):
c = self.c(x)
gm, bt = self.q16(m + ".weight"), self.q16(m + ".bias")
if len(gm) != c or len(bt) != c:
lref()
gb = smallest(gm, (4, 8))
return self.wb.op(0x18, [x], (c,) + A32, struct.pack("<QB7x", 42949, gb) + ints_le(gm, gb) + ints_le(bt, 8))
def TAB(self, x, m):
c = self.c(x)
if c % 4:
lref()
d = c // 4
sa = self.weight(m + ".prep_data.conv1")["shift"]
sf = self.weight(m + ".ffn.project_in")["shift"]
t = self.layernorm(x, m + ".prep_data.norm1")
t = self.conv(t, m + ".prep_data.conv1", 1, 1, 0, 1, sa)
if self.c(t) != 3 * c:
lref()
t = self.conv(t, m + ".prep_data.conv2", 3, 1, 1, "dw", sa)
q, k, v = self.split(t, 0, c), self.split(t, c, c), self.split(t, 2 * c, c)
q = self.unary(0x19, q, struct.pack("<Q", 42949))
k = self.unary(0x19, k, struct.pack("<Q", 42949))
a = self.wb.op(0x1D, [q, k], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
temp = self.q16(m + ".attn.temperature")
if len(temp) != 4:
lref()
mult = [temp[i // d] for i in range(c)]
a = self.wb.op(0x02, [a], (c,) + A32, rq_block(1, 0, mult, 2, smallest(mult, (4, 8)), [16], 1, [], 0))
a = self.unary(0x1A, a)
a = self.wb.op(0x1E, [a, v], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
a = self.conv(a, m + ".attn.project_out", 1, 1, 0, 1, sa)
x2 = self.add(x, a)
t = self.layernorm(x2, m + ".ffn.norm1")
t = self.conv(t, m + ".ffn.project_in", 1, 1, 0, 1, sf)
if self.c(t) % 2:
lref()
h = self.c(t) // 2
t = self.conv(t, m + ".ffn.dwconv", 3, 1, 1, "dw", sf)
a, b = self.split(t, 0, h), self.split(t, h, h)
u = self.mul(self.unary(0x1C, a, struct.pack("<q", 1 << 16)), b)
u = self.conv(u, m + ".ffn.project_out", 1, 1, 0, 1, sf)
return self.add(x2, u)
def rec_graph(hyp, mcms, C):
"""rec_y (mcms: the four context stages' files) or rec_uv (mcms None), 13.6.2."""
n = Q16Net()
n.F = hyp
z = n.wb.input(C, -31, 32, 4, 4)
r = n.wb.input(C, I32_MIN, I32_MAX, 16, 16)
t = n.wb.op(0x02, [z], (C,), rq_block(0, 0, [1 << 16], 1, 4, [0], 1, [], 0))
t = n.conv(t, "conv1", 1, 1, 0, 1)
t = n.convt(t, "conv2", 4, 0)
t = n.relu6(t)
t = n.conv(t, "conv3", 3, 1, 1, 1)
t = n.relu6(t)
psi = n.conv(t, "conv4", 3, 1, 1, 1)
if n.c(psi) != 4 * C:
lref()
if mcms is None:
out = n.add(r, n.merge([n.split(psi, i * C, C) for i in range(4)]))
return n.wb.bytes([out])
rs = n.wb.op_n(0x15, [r], [(C,)] * 4, SUBLATTICE)
ys = []
for i in range(4):
n.F = mcms[i]
p = n.split(psi, i * C, C)
if i == 0:
f = p
else:
ctx = ys[0] if i == 1 else n.concat(ys)
f = n.conv(ctx, "conv.0", 1, 1, 0, 1)
f = n.conv(f, "conv.1", 3, 1, 1, 5)
f = n.concat([f, p])
f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv1", 1, 1, 0, 1))
f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv2", 1, 1, 0, 1))
mu = n.conv(f, "fusion_pred_net.conv3", 1, 1, 0, 1)
if n.c(mu) != C:
lref()
ys.append(n.add(rs[i], mu))
return n.wb.bytes([n.merge(ys)])
def syn_y_graph(F):
n = Q16Net()
n.F = F
x = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
x = n.R(x, "first_stage.0.0")
x = n.convt(x, "first_stage.1", 3, 1)
x = n.G(x, "first_stage.2")
x = n.convt(x, "conv2_t", 3, 1)
x = n.CAB(x)
x = n.G(x, "iact2")
x = n.conv(x, "conv3_t", 1, 1, 0, 1)
x = n.shuffle2(x)
x = n.conv(x, "TAM.ds_conv", 3, 2, 1, 1)
x = n.TAB(x, "TAM.TABs.0")
x = n.TAB(x, "TAM.TABs.1")
x = n.convt(x, "TAM.us_conv", 3, 1)
x = n.G(x, "iact3")
x = n.convt(x, "conv4_t", 3, 1)
if n.c(x) != 1:
lref()
return n.wb.bytes([n.D(x)])
def syn_uv_graph(F):
n = Q16Net()
n.F = F
y = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
uv = n.wb.input(96, I32_MIN, I32_MAX, 16, 16)
x = n.concat([y, uv])
x = n.conv(x, "first_stage.conv1", 3, 1, 1, 1)
x = n.concat([x, uv])
x = n.convt(x, "conv2_t", 3, 1)
x = n.CAB(x)
x = n.G(x, "iact2")
x = n.conv(x, "conv3_t", 1, 1, 0, 1)
x = n.shuffle2(x)
x = n.TAB(x, "TAM.TABs.0")
x = n.TAB(x, "TAM.TABs.1")
x = n.G(x, "iact3")
x = n.convt(x, "conv4_t", 3, 1)
if n.c(x) != 8:
lref()
x = n.shuffle2(x)
return n.wb.bytes([n.D(x)])
def jpegai_scale_lower(norm_b, hy_b, huv_b):
"""jpegai-int/1, section 13.6.1: the scale graphs hs_y and hs_uv."""
try:
norm, hy, huv = json.loads(norm_b), json.loads(hy_b), json.loads(huv_b)
except ValueError:
@ -2554,7 +2924,7 @@ def lower_set(entry, files):
if low == "world/1":
return world_lower(files["T.whm"], files["g_s.whm"], dec(entry["c"]))
if low == "jpegai-int/1":
return jpegai_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
return jpegai_lower(files)
raise ValueError("no lowering %r" % low)
@ -2669,12 +3039,21 @@ def check(root, write, external=None):
return failures
def graph_meets(gj, p):
"""Section 11, `role_profiles`, read from a certificate's graph: its inputs' and
scalars' declared ranges meet profile p's tensor width, and every node fits p."""
tw = WIDTH[p]["tensor"]
ranges = [(dec(r["lo"]), dec(r["hi"])) for r in gj["inputs"] + gj["scalars"]]
return all(bits(r) <= tw for r in ranges) and all(p in nd["fits"] for nd in gj["nodes"])
def check_external(root, external, write):
"""The sets of external.json, whose certificates are committed although the sets
are not. With the set's files (`external` maps an entry's name to the directory
that holds them), the certificate is recomputed and compared. Without them, only
what needs no set: the certificate is canonical, names its set, its capability
and its lowering, and its verdict meets the profile the entry gives."""
and its lowering, its verdict meets the profile the entry gives, and each graph
the entry's `role_profiles` names meets its own profile (section 11)."""
failures = n_full = n_part = 0
for e in json.load(open(os.path.join(root, "certificates", "external.json")))["sets"]:
try:
@ -2696,6 +3075,10 @@ def check_external(root, external, write):
raise ValueError("the certificate names another set")
if c["verdict"][e["profile"]] is not True:
raise ValueError("its verdict does not meet %s" % e["profile"])
for role, rp in sorted(e.get("role_profiles", {}).items()):
gj = [g for g in c["graphs"] if g["name"] == role]
if len(gj) != 1 or not graph_meets(gj[0], rp):
raise ValueError("its graph %s does not meet %s" % (role, rp))
if d:
n_full += 1
else:
@ -2747,23 +3130,32 @@ def check_lower_vectors(root, lowered):
def certify_main(args):
"""`certify`: one set, or WHG1 graphs on their own, to a certificate (section 8.5),
with the exit status of the bounds tool: 0 the verdict meets the profile; 1 it
does not (the certificate is still written); 2 a graph is malformed or refused at
load; 3 `--check` finds a difference; 4 the lowering refuses the set; 5 usage."""
with the exit status of the bounds tool: 0 the verdict meets the profile, and each
graph a --role-profile names meets its own (section 11, `role_profiles`); 1 one of
them does not (the certificate is still written); 2 a graph is malformed or
refused at load; 3 `--check` finds a difference; 4 the lowering refuses the set;
5 usage, a --role-profile role among them that names no graph."""
opts = {"--capability": None, "--lowering": None, "--c": None, "--profile": "acc64", "--out": None, "--check": None}
members, graphs = [], []
members, graphs, role_profiles = [], [], {}
k = 0
while k < len(args):
a = args[k]
if k + 1 >= len(args) or (a not in opts and a not in ("--member", "--whg1")):
if k + 1 >= len(args) or (a not in opts and a not in ("--member", "--whg1", "--role-profile")):
print(CERTIFY_USAGE, file=sys.stderr)
return 5
v = args[k + 1]
if a in ("--member", "--whg1"):
if a in ("--member", "--whg1", "--role-profile"):
if "=" not in v:
print(CERTIFY_USAGE, file=sys.stderr)
return 5
(members if a == "--member" else graphs).append(tuple(v.split("=", 1)))
name, val = v.split("=", 1)
if a == "--role-profile":
if name in role_profiles or val not in PROFILES:
print(CERTIFY_USAGE, file=sys.stderr)
return 5
role_profiles[name] = val
else:
(members if a == "--member" else graphs).append((name, val))
else:
opts[a] = v
k += 2
@ -2786,8 +3178,13 @@ def certify_main(args):
except OSError as ex:
print(ex, file=sys.stderr)
return 5
unknown = sorted(set(role_profiles) - {name for name, _ in roles})
if unknown:
print("--role-profile names no graph of the set: %s" % ", ".join(unknown), file=sys.stderr)
print(CERTIFY_USAGE, file=sys.stderr)
return 5
try:
got, v = certificate(capability, lowering, digest, roles)
got, v, role_fails = certificate(capability, lowering, digest, roles, role_profiles)
except Refuse as ex:
print("a graph is refused at load: %s" % ex.code, file=sys.stderr)
return 2
@ -2802,18 +3199,29 @@ def certify_main(args):
if got != want:
print("--check: %s" % first_difference(got, want), file=sys.stderr)
return 3
status = 0
if not v[p]:
print("the verdict does not meet %s: %s" % (p, cw(v[p + "_first_failure"])), file=sys.stderr)
return 1
status = 1
for name, rp in role_profiles.items():
if name in role_fails:
print("graph %s does not meet %s: %s" % (name, rp, cw(role_fails[name])), file=sys.stderr)
status = 1
else:
print("graph %s meets %s" % (name, rp), file=sys.stderr)
if status:
return status
print("the verdict meets %s: %s" % (p, cw(v)), file=sys.stderr)
return 0
CERTIFY_USAGE = """usage:
wai_ops_check.py certify --capability CAP --lowering ID [--c N] --member NAME=PATH ...
[--profile acc32|acc64] [--out FILE] [--check FILE]
[--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
[--out FILE] [--check FILE]
wai_ops_check.py certify --whg1 ROLE=PATH ... [--capability CAP]
[--profile acc32|acc64] [--out FILE] [--check FILE]"""
[--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
[--out FILE] [--check FILE]"""
def main(argv):

19
wai/wai-ops/src/generate/graphs.rs
View file

@ -537,6 +537,25 @@ pub fn lowered(name: &'static str, set: &str, l: &Lowered) -> File {
let mut r = Rng::named(name);
for (role, gr) in &l.roles {
let b = gr.to_bytes();
// The graphs of `jpegai-int/1` §13.6.2 take more than one input, or none of their
// own size: each decode gives every data input, at a size well below the declared
// one (the outputs are a whole picture's), and one decode puts input 0 one past
// its range.
if matches!(*role, "rec_y" | "rec_uv" | "syn_y" | "syn_uv") {
let sizes: Vec<Vec<(u32, u32)>> = if role.starts_with("rec") { vec![vec![(1, 1), (4, 4)], vec![(2, 1), (8, 4)]] } else { vec![vec![(2, 2); gr.inputs.len()]] };
let inputs = |r: &mut Rng, sz: &[(u32, u32)]| -> Vec<Tensor> {
gr.inputs.iter().zip(sz).map(|(i, &(h, w))| {
let rec = gr.tensors[i.tensor as usize];
tensor(r, rec.channels, h, w, rec.lo.max(-300), rec.hi.min(300))
}).collect()
};
let ds: Vec<Decode> = sizes.iter().map(|sz| Decode::new(inputs(&mut r, sz), vec![])).collect();
g.emit(&format!("{set}/{role}"), &b, ds);
let mut xs = inputs(&mut r, &sizes[0]);
xs[0].data[0] = gr.tensors[gr.inputs[0].tensor as usize].hi + 1;
g.emit(&format!("{set}/{role}/range"), &b, vec![Decode::new(xs, vec![])]);
continue;
}
let i0 = gr.inputs[0];
let rec = gr.tensors[i0.tensor as usize];
let c = rec.channels;

138
wai/wai-ops/src/generate/sets.rs
View file

@ -109,16 +109,21 @@ pub fn synthetic() -> Vec<(Set, Vec<(&'static str, Vec<u8>)>)> {
vec![("T.whm", t), ("g_s.whm", g)],
));
let (norm, hy, huv) = jpegai_bundle(&mut r);
// The integer networks of §13.6.2, from their own sequence, at small widths.
let nets = jpegai_nets(&mut Rng::named("jpegai_nets"));
let mut files: Vec<(&'static str, Vec<u8>)> = vec![("norm.json", norm), ("weights/hs_uv.json", huv), ("weights/hs_y.json", hy)];
files.extend(lower::jpegai_nets::NET_MEMBERS.iter().copied().zip(nets));
files.sort_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
v.push((
Set {
name: "jpegai",
capability: "wai.image.jpegai",
lowering: "jpegai-int/1",
members: vec![("norm.json", base("jpegai", "norm.json")), ("weights/hs_uv.json", base("jpegai", "weights/hs_uv.json")), ("weights/hs_y.json", base("jpegai", "weights/hs_y.json"))],
members: files.iter().map(|(n, _)| (*n, base("jpegai", n))).collect(),
c: None,
note: "synthetic",
},
vec![("norm.json", norm), ("weights/hs_uv.json", huv), ("weights/hs_y.json", hy)],
files,
));
v
}
@ -200,7 +205,14 @@ pub fn lower_set(s: &Set, m: &[(&str, Vec<u8>)]) -> Result<Lowered, String> {
"wam1/1" => lower::wam1(get("model.wam")?).map(|g| Lowered { roles: vec![("g_s", g)] }),
"wmb1/1" => lower::wmb1(get("model.wmb")?),
"world/1" => lower::world(get("T.whm")?, get("g_s.whm")?, s.c.unwrap_or(0)),
"jpegai-int/1" => lower::jpegai(get("norm.json")?, get("weights/hs_y.json")?, get("weights/hs_uv.json")?),
"jpegai-int/1" => {
let names = ["norm.json", "weights/hs_y.json", "weights/hs_uv.json"].into_iter().chain(lower::jpegai_nets::NET_MEMBERS);
for n in names {
get(n)?;
}
let b = lower::JpegaiBytes::from_members(|n| m.iter().find(|(k, _)| *k == n).map(|(_, b)| b.as_slice())).ok_or_else(|| format!("{}: a member is missing", s.name))?;
lower::jpegai(&b)
}
_ => return Err(format!("{}: no lowering {}", s.name, s.lowering)),
};
r.map_err(|e| format!("{}: the lowering refused ({:?})", s.name, e))
@ -497,6 +509,126 @@ fn jpegai_bundle(r: &mut Rng) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
(norm.into_bytes(), hy, huv)
}
/// The eight weight files of §13.6.2 at small widths, in
/// [`lower::jpegai_nets::NET_MEMBERS`]' order: hyper decoders of width 2, context stages
/// of width 2 (their first convolutions 5 wide, for the groups of 5), and synthesis
/// networks of width 4 (four heads of one channel each). Each convolution has its own
/// shift, so the TAB rule (each half takes its first convolution's) is visible in the
/// lowered bytes, and a multiplier that keeps its input's magnitude, so that the graphs'
/// vectors decode through every operator rather than refuse past `A32`.
fn jpegai_nets(r: &mut Rng) -> Vec<Vec<u8>> {
struct F(Vec<String>);
impl F {
/// `m.weight`, its multiplier set so that a weight of 32767 times its `taps`
/// products keeps the input's magnitude: the graphs' decodes complete.
fn weight(&mut self, r: &mut Rng, name: &str, dims: [usize; 4], taps: usize) {
let n: usize = dims.iter().product();
let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
let shift = r.int(26, 31);
let mult = ((1i64 << shift) / (32767 * taps as i64)).max(1) * r.int(1, 2);
self.0.push(format!("\"{name}.weight\":{{\"kind\":\"conv\",\"dims\":{},\"mult\":{},\"shift\":{},\"w\":{}}}", ints(&d), mult, shift, ints(&rv(r, n, 32767))));
}
fn conv(&mut self, r: &mut Rng, name: &str, dims: [usize; 4], bias: bool) {
self.weight(r, name, dims, dims[1] * dims[2] * dims[3]);
if bias {
self.q16(&format!("{name}.bias"), &rv(r, dims[0], 1 << 16));
}
}
fn convt(&mut self, r: &mut Rng, name: &str, dims: [usize; 4]) {
self.weight(r, name, dims, dims[0] * dims[2] * dims[3]);
self.q16(&format!("{name}.bias"), &rv(r, dims[1], 1 << 16));
}
fn q16(&mut self, name: &str, v: &[i64]) {
self.0.push(format!("\"{name}\":{{\"kind\":\"q16\",\"dims\":[{}],\"v\":{}}}", v.len(), ints(v)));
}
fn res(&mut self, r: &mut Rng, m: &str, c: usize, mid: usize) {
self.conv(r, &format!("{m}.conv1"), [mid, c, 3, 3], true);
self.conv(r, &format!("{m}.conv2"), [c, mid, 3, 3], true);
}
fn gate(&mut self, r: &mut Rng, m: &str, c: usize) {
self.conv(r, &format!("{m}.conv"), [c, c / 4, 3, 3], false);
self.conv(r, &format!("{m}.conv2"), [c, c, 1, 1], false);
}
fn cab(&mut self, r: &mut Rng, c: usize) {
self.res(r, "CAB.residual_trunk.0", c, 2);
self.res(r, "CAB.residual_trunk.1", c, 2);
self.conv(r, "CAB.subscale", [c, c, 3, 3], true);
self.res(r, "CAB.residual_mask2.0", c, 2);
self.res(r, "CAB.residual_mask3.0", c, 2);
self.convt(r, "CAB.upscale", [c, c, 3, 3]);
}
fn tab(&mut self, r: &mut Rng, m: &str, c: usize, hid: usize) {
let pos = |r: &mut Rng, n: usize| -> Vec<i64> { (0..n).map(|_| r.int(0, 1 << 17)).collect() };
let g = pos(r, c);
self.q16(&format!("{m}.prep_data.norm1.weight"), &g);
self.q16(&format!("{m}.prep_data.norm1.bias"), &rv(r, c, 1 << 15));
self.conv(r, &format!("{m}.prep_data.conv1"), [3 * c, c, 1, 1], false);
self.conv(r, &format!("{m}.prep_data.conv2"), [3 * c, 1, 3, 3], false);
let t: Vec<i64> = (0..4).map(|_| r.int(1 << 16, 1 << 19)).collect();
self.q16(&format!("{m}.attn.temperature"), &t);
self.conv(r, &format!("{m}.attn.project_out"), [c, c, 1, 1], false);
let g = pos(r, c);
self.q16(&format!("{m}.ffn.norm1.weight"), &g);
self.q16(&format!("{m}.ffn.norm1.bias"), &rv(r, c, 1 << 15));
self.conv(r, &format!("{m}.ffn.project_in"), [2 * hid, c, 1, 1], false);
self.conv(r, &format!("{m}.ffn.dwconv"), [2 * hid, 1, 3, 3], false);
self.conv(r, &format!("{m}.ffn.project_out"), [c, hid, 1, 1], false);
}
fn done(self) -> Vec<u8> {
format!("{{{}}}\n", self.0.join(",")).into_bytes()
}
}
let mut out = Vec::new();
for cz in [160usize, 96] {
let mut f = F(Vec::new());
f.conv(r, "conv1", [2, cz, 1, 1], false);
f.convt(r, "conv2", [2, 2, 4, 4]);
f.conv(r, "conv3", [2, 2, 3, 3], true);
f.conv(r, "conv4", [4 * cz, 2, 3, 3], true);
out.push(f.done());
}
for i in 0..4usize {
let mut f = F(Vec::new());
let cin = if i == 0 { 160 } else {
f.conv(r, "conv.0", [5, 160 * i, 1, 1], false);
f.conv(r, "conv.1", [5, 1, 3, 3], true);
165
};
f.conv(r, "fusion_pred_net.conv1", [2, cin, 1, 1], false);
f.conv(r, "fusion_pred_net.conv2", [2, 2, 1, 1], false);
f.conv(r, "fusion_pred_net.conv3", [160, 2, 1, 1], false);
out.push(f.done());
}
let w = 4usize;
let mut f = F(Vec::new());
f.res(r, "first_stage.0.0", 160, 2);
f.convt(r, "first_stage.1", [160, w, 3, 3]);
f.gate(r, "first_stage.2", w);
f.convt(r, "conv2_t", [w, w, 3, 3]);
f.cab(r, w);
f.gate(r, "iact2", w);
f.conv(r, "conv3_t", [4 * w, w, 1, 1], true);
f.conv(r, "TAM.ds_conv", [w, w, 3, 3], true);
f.tab(r, "TAM.TABs.0", w, 2);
f.tab(r, "TAM.TABs.1", w, 2);
f.convt(r, "TAM.us_conv", [w, w, 3, 3]);
f.gate(r, "iact3", w);
f.convt(r, "conv4_t", [w, 1, 3, 3]);
out.push(f.done());
let mut f = F(Vec::new());
f.conv(r, "first_stage.conv1", [2, 256, 3, 3], true);
f.convt(r, "conv2_t", [98, w, 3, 3]);
f.cab(r, w);
f.gate(r, "iact2", w);
f.conv(r, "conv3_t", [4 * w, w, 1, 1], true);
f.tab(r, "TAM.TABs.0", w, 2);
f.tab(r, "TAM.TABs.1", w, 2);
f.gate(r, "iact3", w);
f.convt(r, "conv4_t", [w, 8, 3, 3]);
out.push(f.done());
out
}
/// Write the index (`certificates/index.json`): pretty enough to read, and
/// canonical in its member order.
/// The entry of a falsification set: an index entry with the profile it is checked

50
wai/wai-ops/src/lower/jpegai.rs
View file

@ -1,10 +1,12 @@
//! `jpegai-int/1` (§13.6.1): the scale graphs `hs_y` and `hs_uv` of
//! `wai.image.jpegai`, from the bundle's `norm.json` and `weights/hs_{y,uv}.json`.
//! `jpegai-int/1`: the integer stages of `wai.image.jpegai`. This module lowers the
//! scale path (§13.6.1), the graphs `hs_y` and `hs_uv`, from the bundle's `norm.json`
//! and `weights/hs_{y,uv}.json`; [`super::jpegai_nets`] lowers the other integer
//! networks (§13.6.2), and [`jpegai`] writes the six graphs in §13.0's role order.
//!
//! [`jpegai`] lowers the members' bytes. A decoder that holds the members parsed lowers
//! them from that form ([`jpegai_net`], [`jpegai_net_loaded`]), which is
//! [`parse_jpegai`] of those bytes: the graphs are the ones [`jpegai`] writes, and the
//! refusals are the same.
//! [`jpegai_scale`] lowers the scale path's members' bytes. A decoder that holds the
//! members parsed lowers them from that form ([`jpegai_net`], [`jpegai_net_loaded`]),
//! which is [`parse_jpegai`] of those bytes: the graphs are the ones [`jpegai_scale`]
//! writes, and the refusals are the same.
use alloc::vec;
use alloc::vec::Vec;
@ -222,11 +224,45 @@ fn graphs(net: JpegaiNet) -> Res<(Graph, Graph)> {
}
/// Lower a `wai.image.jpegai` bundle's scale path: `hs_y`, then `hs_uv`.
pub fn jpegai(norm: &[u8], hs_y: &[u8], hs_uv: &[u8]) -> Res<Lowered> {
pub fn jpegai_scale(norm: &[u8], hs_y: &[u8], hs_uv: &[u8]) -> Res<Lowered> {
let (y, uv) = graphs(parse_jpegai(norm, hs_y, hs_uv)?)?;
Ok(Lowered { roles: vec![("hs_y", y), ("hs_uv", uv)] })
}
/// The members of a `wai.image.jpegai` bundle that `jpegai-int/1` reads, as bytes:
/// `norm.json`, `weights/hs_y.json`, `weights/hs_uv.json`, then the eight files of
/// [`super::jpegai_nets::NET_MEMBERS`] in that order.
#[derive(Clone, Copy, Debug)]
pub struct JpegaiBytes<'a> {
pub norm: &'a [u8],
pub hs_y: &'a [u8],
pub hs_uv: &'a [u8],
pub nets: [&'a [u8]; 8],
}
impl<'a> JpegaiBytes<'a> {
/// Every member the lowering reads, from `get` by its name in the set
/// (`norm.json`, `weights/…`); `None` when one is missing.
pub fn from_members(get: impl Fn(&str) -> Option<&'a [u8]>) -> Option<JpegaiBytes<'a>> {
let n = super::jpegai_nets::NET_MEMBERS;
Some(JpegaiBytes {
norm: get("norm.json")?,
hs_y: get("weights/hs_y.json")?,
hs_uv: get("weights/hs_uv.json")?,
nets: [get(n[0])?, get(n[1])?, get(n[2])?, get(n[3])?, get(n[4])?, get(n[5])?, get(n[6])?, get(n[7])?],
})
}
}
/// Lower a `wai.image.jpegai` bundle: the scale path `hs_y` and `hs_uv` (§13.6.1), then
/// `rec_y`, `rec_uv`, `syn_y` and `syn_uv` (§13.6.2).
pub fn jpegai(b: &JpegaiBytes<'_>) -> Res<Lowered> {
let mut l = jpegai_scale(b.norm, b.hs_y, b.hs_uv)?;
let nets = super::jpegai_nets::nets_graphs(super::jpegai_nets::parse_nets(b.nets)?)?;
l.roles.extend(super::jpegai_nets::NET_ROLES.iter().copied().zip(nets));
Ok(l)
}
/// [`jpegai`] of a bundle already parsed ([`parse_jpegai`] of its bytes, or the same
/// members from another parser): the graphs [`jpegai`] writes from those bytes, and
/// its refusals.

550
wai/wai-ops/src/lower/jpegai_nets.rs Normal file
View file

@ -0,0 +1,550 @@
//! `jpegai-int/1` (§13.6.2): the other integer networks of `wai.image.jpegai`, the
//! hyper decoder, the four context stages and the synthesis, as the graphs `rec_y`,
//! `rec_uv`, `syn_y` and `syn_uv`, from the bundle's `weights/hyp_{y,uv}.json`,
//! `weights/mcm{0,1,2,3}_y.json` and `weights/synth_{y,uv}.json`.
//!
//! [`parse_ptq`] reads one weight file's members (§13.6.2, *The members*). A decoder that
//! holds the files parsed another way converts them to [`PtqNet`] and lowers that, which
//! is the lowering of the bytes. Every tensor these graphs write that an operator can
//! grow (the output of a convolution, a transposed convolution, an addition, a product,
//! a layer normalisation, an attention contraction or a scale) has the policy
//! `check [−2^31, 2^31 − 1]`, and so does every input but the hyper-latent.
use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::String;
use alloc::vec;
use alloc::vec::Vec;
use super::{Builder, Lowered, Res, refuse};
use crate::bounds::Profile;
use crate::graph::{BiasKind, Graph, Loaded, Params, Round, Rq, TensorRec, opcode, smallest_width};
use crate::json::{self, Value};
/// The range every checked activation of these graphs keeps to: it fits i32.
pub const A32: (i64, i64) = (i32::MIN as i64, i32::MAX as i64);
/// `1.0` in Q16.
pub const ONE: i64 = 1 << 16;
/// ReLU6's upper end, `6·2^16`.
pub const RELU6: i64 = 6 << 16;
/// The `ε` of every layer and L2 normalisation, `⌊10^−5·2^32⌋`.
pub const EPS: u64 = 42949;
/// The heads of every channel attention.
pub const HEADS: u32 = 4;
/// The denormalisation's bias, `127·2^16 + 2^15`.
pub const DENORM_BIAS: i64 = 127 * ONE + ONE / 2;
/// The denormalisation's output range, `[0, 255·2^16]`.
pub const DENORM_HI: i64 = 255 * ONE;
/// The 2×2 sub-lattices in stage order, `(0,0), (1,1), (0,1), (1,0)`, each `2a + b`.
pub const SUBPOS: [u8; 4] = [0, 3, 1, 2];
/// The latent's side and the hyper-latent's.
pub const SIDE: u32 = 16;
pub const Z_SIDE: u32 = 4;
/// The latent channels of each plane, luma then chroma.
pub const CHANNELS: [u32; 2] = [160, 96];
/// One member of a weight file that §13.6.2 reads, as parsed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PtqMember {
/// Kind `conv`: the weights' `dims`, the weights `w` row-major over them, the one
/// multiplier and the one shift.
Conv { dims: Vec<u64>, w: Vec<i64>, mult: i64, shift: u8 },
/// Kind `q16`: the values.
Q16 { v: Vec<i64> },
}
/// A weight file as parsed: its members of kind `conv` and `q16`, by name. Members of
/// another kind are not kept.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct PtqNet {
pub members: BTreeMap<String, PtqMember>,
}
impl PtqNet {
/// Add a member, refusing a second one of the same name.
pub fn insert(&mut self, name: String, m: PtqMember) -> Res<()> {
match self.members.insert(name, m) {
None => Ok(()),
Some(_) => refuse(),
}
}
}
/// The integer networks of a bundle as parsed: the two hyper decoders, the four
/// context stages (luma only) and the two synthesis networks.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct JpegaiNets {
pub hyp: [PtqNet; 2],
pub mcm: [PtqNet; 4],
pub synth: [PtqNet; 2],
}
/// The members' file names, in [`JpegaiNets`]' order.
pub const NET_MEMBERS: [&str; 8] = [
"weights/hyp_y.json",
"weights/hyp_uv.json",
"weights/mcm0_y.json",
"weights/mcm1_y.json",
"weights/mcm2_y.json",
"weights/mcm3_y.json",
"weights/synth_y.json",
"weights/synth_uv.json",
];
fn ints(v: &Value) -> Res<Vec<i64>> {
let a = v.as_arr().ok_or(()).or_else(|_| refuse())?;
a.iter().map(|x| x.as_i64().ok_or(()).or_else(|_| refuse())).collect()
}
/// Whether some object of `v`, at any depth, has two members of one name.
fn duplicate_names(v: &Value) -> bool {
match v {
Value::Obj(m) => {
let mut names: Vec<&str> = m.iter().map(|(k, _)| k.as_str()).collect();
names.sort_unstable();
names.windows(2).any(|p| p[0] == p[1]) || m.iter().any(|(_, x)| duplicate_names(x))
}
Value::Arr(a) => a.iter().any(duplicate_names),
_ => false,
}
}
/// One weight file's members (§13.6.2, *The members*): a JSON object, each member of
/// kind `conv` with integer `dims`, as many weights `w` as they give (each fitting i64),
/// an integer `mult` and a `shift` in `0..=127`, each of kind `q16` with integer values
/// `v`; and no object, at any depth, with two members of one name.
pub fn parse_ptq(b: &[u8]) -> Res<PtqNet> {
let v = json::parse(b).or_else(|_| refuse())?;
let obj = v.as_obj().ok_or(()).or_else(|_| refuse())?;
if duplicate_names(&v) {
return refuse();
}
let mut net = PtqNet::default();
for (name, e) in obj {
let kind = e.as_obj().and_then(|_| e.get("kind")).and_then(Value::as_str);
let get = |k: &str| e.get(k).ok_or(()).or_else(|_| refuse());
let m = match kind {
Some("conv") => {
let dims: Vec<u64> = get("dims")?.as_arr().ok_or(()).or_else(|_| refuse())?.iter().map(|d| d.as_u64().ok_or(()).or_else(|_| refuse())).collect::<Res<_>>()?;
let shift = get("shift")?.as_u64().filter(|&s| s <= 127).ok_or(()).or_else(|_| refuse())? as u8;
let w = ints(get("w")?)?;
if product(&dims) != Some(w.len() as u64) {
return refuse();
}
PtqMember::Conv { dims, w, mult: get("mult")?.as_i64().ok_or(()).or_else(|_| refuse())?, shift }
}
Some("q16") => PtqMember::Q16 { v: ints(get("v")?)? },
_ => continue,
};
net.insert(name.clone(), m)?;
}
Ok(net)
}
/// The eight files, parsed.
pub fn parse_nets(files: [&[u8]; 8]) -> Res<JpegaiNets> {
let [hy, huv, m0, m1, m2, m3, sy, suv] = files;
Ok(JpegaiNets { hyp: [parse_ptq(hy)?, parse_ptq(huv)?], mcm: [parse_ptq(m0)?, parse_ptq(m1)?, parse_ptq(m2)?, parse_ptq(m3)?], synth: [parse_ptq(sy)?, parse_ptq(suv)?] })
}
// ------------------------------------------------------------------ building blocks
/// A graph under construction over one weight file at a time.
struct G {
b: Builder,
}
/// A convolution's geometry as the topology fixes it: kernel, stride, padding, and
/// groups (`0` for depthwise, `groups = cin`).
#[derive(Clone, Copy)]
struct Geo {
k: u32,
s: u32,
p: u32,
g: u32,
}
const fn geo(k: u32, s: u32, p: u32, g: u32) -> Geo {
Geo { k, s, p, g }
}
/// 1×1, stride 1, padding 0.
const P1: Geo = geo(1, 1, 0, 1);
/// 3×3, stride 1, padding 1.
const C3: Geo = geo(3, 1, 1, 1);
/// 3×3, stride 2, padding 1: the downsamplers, and (with output padding 1) the
/// transposed upsamplers.
const S2: Geo = geo(3, 2, 1, 1);
fn checked(c: u32) -> TensorRec {
TensorRec::check(c, A32.0, A32.1)
}
fn take_conv(net: &mut PtqNet, m: &str) -> Res<(Vec<u64>, Vec<i64>, i64, u8)> {
match net.members.get_mut(&format!("{m}.weight")) {
Some(PtqMember::Conv { dims, w, mult, shift }) => Ok((dims.clone(), core::mem::take(w), *mult, *shift)),
_ => refuse(),
}
}
/// The shift of `m.weight`, which must be a `conv` member.
fn shift_of(net: &PtqNet, m: &str) -> Res<u8> {
match net.members.get(&format!("{m}.weight")) {
Some(PtqMember::Conv { shift, .. }) => Ok(*shift),
_ => refuse(),
}
}
/// `name`, a `q16` member of exactly `n` values.
fn take_q16(net: &mut PtqNet, name: &str, n: u64) -> Res<Vec<i64>> {
match net.members.get_mut(name) {
Some(PtqMember::Q16 { v }) if v.len() as u64 == n => Ok(core::mem::take(v)),
_ => refuse(),
}
}
/// `m.bias`: `None` when the file has no such member; otherwise a `q16` member of
/// `cout` values.
fn take_bias(net: &mut PtqNet, m: &str, cout: u64) -> Res<Option<Vec<i64>>> {
let name = format!("{m}.bias");
if !net.members.contains_key(&name) {
return Ok(None);
}
take_q16(net, &name, cout).map(Some)
}
fn product(d: &[u64]) -> Option<u64> {
d.iter().try_fold(1u64, |a, &x| a.checked_mul(x))
}
/// The convolutions' `RQ`: one multiplier, one shift, a `post` bias when the file has
/// one, `floor`.
fn conv_rq(mult: i64, shift: u8, bias: Option<Vec<i64>>) -> Rq {
let (bias_kind, b_bytes, bias) = match bias {
Some(b) => (BiasKind::Post, smallest_width(&b, &[4, 8]), b),
None => (BiasKind::None, 0, vec![]),
};
Rq { round: Round::Floor, bias_kind, m_kind: 1, s_kind: 1, m_bytes: smallest_width(&[mult], &[4, 8]), b_bytes, mult: vec![mult], shift: vec![shift], bias }
}
impl G {
fn new() -> G {
G { b: Builder::new() }
}
fn ch(&self, t: u32) -> u32 {
self.b.channels(t)
}
/// `conv` from `m.weight` (`[cout, cin/groups, k, k]`), output `check`; `shift`, when
/// given, in place of the member's own (§13.6.2, the TAB rule).
fn conv(&mut self, net: &mut PtqNet, m: &str, x: u32, ge: Geo, shift: Option<u8>) -> Res<u32> {
let cin = self.ch(x) as u64;
let groups = if ge.g == 0 { cin } else { ge.g as u64 };
let (dims, w, mult, own) = take_conv(net, m)?;
let k = ge.k as u64;
if dims.len() != 4 || groups == 0 || cin % groups != 0 || dims[1].checked_mul(groups) != Some(cin) || dims[2] != k || dims[3] != k {
return refuse();
}
let cout = dims[0];
if cout == 0 || cout > u32::MAX as u64 || cout % groups != 0 || (ge.g == 0 && cout != cin) || product(&dims) != Some(w.len() as u64) {
return refuse();
}
let bias = take_bias(net, m, cout)?;
let rq = conv_rq(mult, shift.unwrap_or(own), bias);
let w_bytes = smallest_width(&w, &[1, 2, 4, 8]);
let p = Params::Conv { cout: cout as u32, kh: ge.k, kw: ge.k, sh: ge.s, sw: ge.s, ph: ge.p, pw: ge.p, groups: groups as u32, w_bytes, rq, w };
Ok(self.b.op(opcode::CONV, vec![x], checked(cout as u32), p))
}
/// `conv_transpose` from `m.weight` (`[cin, cout, k, k]`), output padding `op`,
/// output `check`.
fn convt(&mut self, net: &mut PtqNet, m: &str, x: u32, ge: Geo, op: u32) -> Res<u32> {
let cin = self.ch(x) as u64;
let (dims, w, mult, shift) = take_conv(net, m)?;
let k = ge.k as u64;
if dims.len() != 4 || dims[0] != cin || dims[2] != k || dims[3] != k {
return refuse();
}
let cout = dims[1];
if cout == 0 || cout > u32::MAX as u64 || product(&dims) != Some(w.len() as u64) {
return refuse();
}
let bias = take_bias(net, m, cout)?;
let rq = conv_rq(mult, shift, bias);
let w_bytes = smallest_width(&w, &[1, 2, 4, 8]);
let p = Params::ConvT { cout: cout as u32, kh: ge.k, kw: ge.k, sh: ge.s, sw: ge.s, ph: ge.p, pw: ge.p, oph: op, opw: op, w_bytes, rq, w };
Ok(self.b.op(opcode::CONV_TRANSPOSE, vec![x], checked(cout as u32), p))
}
fn relu(&mut self, x: u32) -> u32 {
let c = self.ch(x);
self.b.op(opcode::RELU, vec![x], TensorRec::derived(c), Params::None)
}
fn relu6(&mut self, x: u32) -> u32 {
let c = self.ch(x);
self.b.op(opcode::CLAMP, vec![x], TensorRec::derived(c), Params::Clamp { lo: 0, hi: RELU6 })
}
fn add(&mut self, a: u32, c: u32) -> Res<u32> {
let n = self.ch(a);
if self.ch(c) != n {
return refuse();
}
Ok(self.b.op(opcode::ADD, vec![a, c], checked(n), Params::None))
}
/// A Q16 product: `mul`, no multiplier, shift 16, `floor`.
fn mul(&mut self, a: u32, c: u32) -> Res<u32> {
let n = self.ch(a);
if self.ch(c) != n {
return refuse();
}
Ok(self.b.op(opcode::MUL, vec![a, c], checked(n), Params::Mul(Rq::plain(Round::Floor, 16))))
}
fn split(&mut self, x: u32, start: u32, count: u32) -> u32 {
self.b.op(opcode::SPLIT, vec![x], TensorRec::derived(count), Params::Split { start, count })
}
fn concat(&mut self, xs: Vec<u32>) -> u32 {
let c = xs.iter().map(|&t| self.ch(t)).sum();
self.b.op(opcode::CONCAT, xs, TensorRec::derived(c), Params::None)
}
fn shuffle(&mut self, x: u32) -> Res<u32> {
let c = self.ch(x);
if c % 4 != 0 {
return refuse();
}
Ok(self.b.op(opcode::PIXEL_SHUFFLE, vec![x], TensorRec::derived(c / 4), Params::Shuffle { r: 2 }))
}
/// The residual block `m`: `conv1` (3×3), `relu`, `conv2` (3×3), then `add(x, ·)`.
fn res(&mut self, net: &mut PtqNet, m: &str, x: u32) -> Res<u32> {
let a = self.conv(net, &format!("{m}.conv1"), x, C3, None)?;
let a = self.relu(a);
let b = self.conv(net, &format!("{m}.conv2"), a, C3, None)?;
self.add(x, b)
}
/// The gate `m`: `x · (f(x) + 1)` in Q16, where `f` is ReLU6, `m.conv` (3×3, groups
/// 4) and `m.conv2` (1×1).
fn gate(&mut self, net: &mut PtqNet, m: &str, x: u32) -> Res<u32> {
let a = self.relu6(x);
let a = self.conv(net, &format!("{m}.conv"), a, geo(3, 1, 1, 4), None)?;
let a = self.conv(net, &format!("{m}.conv2"), a, P1, None)?;
let c = self.ch(a);
let one = Rq { round: Round::Floor, bias_kind: BiasKind::Post, m_kind: 0, s_kind: 1, m_bytes: 0, b_bytes: 4, mult: vec![], shift: vec![0], bias: vec![ONE; c as usize] };
let a = self.b.op(opcode::RESCALE, vec![a], checked(c), Params::Rescale(one));
self.mul(x, a)
}
/// `layernorm_c` from `m.weight` (`γ`) and `m.bias` (`β`), each of `c` values.
fn layernorm(&mut self, net: &mut PtqNet, m: &str, x: u32) -> Res<u32> {
let c = self.ch(x);
let gamma = take_q16(net, &format!("{m}.weight"), c as u64)?;
let beta = take_q16(net, &format!("{m}.bias"), c as u64)?;
let g_bytes = smallest_width(&gamma, &[4, 8]);
Ok(self.b.op(opcode::LAYERNORM_C, vec![x], checked(c), Params::LayerNorm { eps: EPS, g_bytes, gamma, beta }))
}
/// One TAB `m` on `x` (`c` channels): the channel-attention half, then the
/// feed-forward half, each a residual.
fn tab(&mut self, net: &mut PtqNet, m: &str, x: u32) -> Res<u32> {
let c = self.ch(x);
if c % HEADS != 0 {
return refuse();
}
let d = c / HEADS;
// The attention half: every convolution takes `prep_data.conv1`'s shift.
let sa = shift_of(net, &format!("{m}.prep_data.conv1"))?;
let y = self.layernorm(net, &format!("{m}.prep_data.norm1"), x)?;
let qkv = self.conv(net, &format!("{m}.prep_data.conv1"), y, P1, Some(sa))?;
if self.ch(qkv) != 3 * c {
return refuse();
}
let qkv = self.conv(net, &format!("{m}.prep_data.conv2"), qkv, geo(3, 1, 1, 0), Some(sa))?;
let q = self.split(qkv, 0, c);
let k = self.split(qkv, c, c);
let v = self.split(qkv, 2 * c, c);
let q = self.b.op(opcode::L2NORM, vec![q], TensorRec::derived(c), Params::L2norm { eps: EPS });
let k = self.b.op(opcode::L2NORM, vec![k], TensorRec::derived(c), Params::L2norm { eps: EPS });
let s = self.b.op(opcode::GRAM, vec![q, k], checked(c), Params::Gram { groups: HEADS, rq: Rq::plain(Round::HalfUp, 16) });
let temp = take_q16(net, &format!("{m}.attn.temperature"), HEADS as u64)?;
let mult: Vec<i64> = (0..c).map(|ch| temp[(ch / d) as usize]).collect();
let m_bytes = smallest_width(&mult, &[4, 8]);
let rq = Rq { round: Round::HalfUp, bias_kind: BiasKind::None, m_kind: 2, s_kind: 1, m_bytes, b_bytes: 0, mult, shift: vec![16], bias: vec![] };
let s = self.b.op(opcode::RESCALE, vec![s], checked(c), Params::Rescale(rq));
let a = self.b.op(opcode::SOFTMAX, vec![s], TensorRec::derived(c), Params::None);
let o = self.b.op(opcode::APPLY, vec![a, v], checked(c), Params::Apply { groups: HEADS, rq: Rq::plain(Round::HalfUp, 16) });
let o = self.conv(net, &format!("{m}.attn.project_out"), o, P1, Some(sa))?;
let x = self.add(x, o)?;
// The feed-forward half: every convolution takes `ffn.project_in`'s shift.
let sf = shift_of(net, &format!("{m}.ffn.project_in"))?;
let y = self.layernorm(net, &format!("{m}.ffn.norm1"), x)?;
let h2 = self.conv(net, &format!("{m}.ffn.project_in"), y, P1, Some(sf))?;
let n2 = self.ch(h2);
if n2 % 2 != 0 {
return refuse();
}
let h2 = self.conv(net, &format!("{m}.ffn.dwconv"), h2, geo(3, 1, 1, 0), Some(sf))?;
let a = self.split(h2, 0, n2 / 2);
let b = self.split(h2, n2 / 2, n2 / 2);
let a = self.b.op(opcode::ELU, vec![a], TensorRec::derived(n2 / 2), Params::Elu { alpha: ONE });
let g = self.mul(a, b)?;
let o = self.conv(net, &format!("{m}.ffn.project_out"), g, P1, Some(sf))?;
self.add(x, o)
}
/// The channel-attention block `m` of both synthesis networks: the trunk (two
/// residual blocks) times the sigmoid of the mask branch, plus `x`.
fn cab(&mut self, net: &mut PtqNet, x: u32) -> Res<u32> {
let t = self.res(net, "CAB.residual_trunk.0", x)?;
let t = self.res(net, "CAB.residual_trunk.1", t)?;
let s = self.conv(net, "CAB.subscale", x, S2, None)?;
let s = self.res(net, "CAB.residual_mask2.0", s)?;
let s = self.res(net, "CAB.residual_mask3.0", s)?;
let u = self.convt(net, "CAB.upscale", s, S2, 1)?;
let c = self.ch(u);
let mask = self.b.op(opcode::SIGMOID, vec![u], TensorRec::derived(c), Params::None);
let p = self.mul(t, mask)?;
self.add(x, p)
}
/// The denormalisation: `rescale` with one multiplier 128, shift 0 and a `post` bias
/// of `127·2^16 + 2^15`, its output `clamp [0, 255·2^16]`.
fn denorm(&mut self, x: u32) -> u32 {
let c = self.ch(x);
let rq = Rq { round: Round::Floor, bias_kind: BiasKind::Post, m_kind: 1, s_kind: 1, m_bytes: 4, b_bytes: 4, mult: vec![128], shift: vec![0], bias: vec![DENORM_BIAS; c as usize] };
self.b.op(opcode::RESCALE, vec![x], TensorRec::clamp(c, 0, DENORM_HI), Params::Rescale(rq))
}
}
// ------------------------------------------------------------------ the graphs
/// A plane's graph `rec_y` (`mcm` given) or `rec_uv`: inputs `z [C, 4, 4]` and the
/// dequantised residual `r [C, 16, 16]`; output `ŷ [C, 16, 16]`.
pub fn rec(hyp: &mut PtqNet, mcm: Option<&mut [PtqNet; 4]>, chroma: bool) -> Res<Graph> {
let cz = CHANNELS[chroma as usize];
let mut g = G::new();
let z = g.b.input(TensorRec::check(cz, -31, 32), Z_SIDE, Z_SIDE);
let r = g.b.input(checked(cz), SIDE, SIDE);
// The hyper decoder ψ on `z·2^16`.
let lift = Rq { round: Round::Floor, bias_kind: BiasKind::None, m_kind: 1, s_kind: 1, m_bytes: 4, b_bytes: 0, mult: vec![ONE], shift: vec![0], bias: vec![] };
let x = g.b.op(opcode::RESCALE, vec![z], TensorRec::derived(cz), Params::Rescale(lift));
let x = g.conv(hyp, "conv1", x, P1, None)?;
let x = g.convt(hyp, "conv2", x, geo(4, 2, 1, 1), 0)?;
let x = g.relu6(x);
let x = g.conv(hyp, "conv3", x, C3, None)?;
let x = g.relu6(x);
let psi = g.conv(hyp, "conv4", x, C3, None)?;
if g.ch(psi) != 4 * cz {
return refuse();
}
let out = match mcm {
None => {
// Chroma: the hyper decoder's four parts are the means.
let parts: Vec<u32> = (0..4).map(|i| g.split(psi, i * cz, cz)).collect();
let mu = g.b.op(opcode::SUBPOS_MERGE, parts, TensorRec::derived(cz), Params::Subpos { order: SUBPOS });
g.add(r, mu)?
}
Some(mcm) => {
let rs = g.b.ops(opcode::SUBPOS_SPLIT, vec![r], vec![TensorRec::derived(cz); 4], Params::Subpos { order: SUBPOS });
let mut ys: Vec<u32> = Vec::with_capacity(4);
for (i, st) in mcm.iter_mut().enumerate() {
let p = g.split(psi, i as u32 * cz, cz);
let f = if i == 0 {
p
} else {
let ctx = if i == 1 { ys[0] } else { g.concat(ys.clone()) };
let c0 = g.conv(st, "conv.0", ctx, P1, None)?;
let cy = g.conv(st, "conv.1", c0, geo(3, 1, 1, 5), None)?;
g.concat(vec![cy, p])
};
let f = g.conv(st, "fusion_pred_net.conv1", f, P1, None)?;
let f = g.relu(f);
let f = g.conv(st, "fusion_pred_net.conv2", f, P1, None)?;
let f = g.relu(f);
let mu = g.conv(st, "fusion_pred_net.conv3", f, P1, None)?;
if g.ch(mu) != cz {
return refuse();
}
ys.push(g.add(rs[i], mu)?);
}
g.b.op(opcode::SUBPOS_MERGE, ys, TensorRec::derived(cz), Params::Subpos { order: SUBPOS })
}
};
g.b.finish(vec![out])
}
/// The luma synthesis `syn_y`: input the luma latent `[160, 16, 16]`; output the
/// denormalised plane `[1, 256, 256]`.
pub fn syn_y(net: &mut PtqNet) -> Res<Graph> {
let mut g = G::new();
let l = g.b.input(checked(CHANNELS[0]), SIDE, SIDE);
let x = g.res(net, "first_stage.0.0", l)?;
let x = g.convt(net, "first_stage.1", x, S2, 1)?;
let x = g.gate(net, "first_stage.2", x)?;
let x = g.convt(net, "conv2_t", x, S2, 1)?;
let x = g.cab(net, x)?;
let x = g.gate(net, "iact2", x)?;
let x = g.conv(net, "conv3_t", x, P1, None)?;
let x = g.shuffle(x)?;
let x = g.conv(net, "TAM.ds_conv", x, S2, None)?;
let x = g.tab(net, "TAM.TABs.0", x)?;
let x = g.tab(net, "TAM.TABs.1", x)?;
let x = g.convt(net, "TAM.us_conv", x, S2, 1)?;
let x = g.gate(net, "iact3", x)?;
let x = g.convt(net, "conv4_t", x, S2, 1)?;
if g.ch(x) != 1 {
return refuse();
}
let x = g.denorm(x);
g.b.finish(vec![x])
}
/// The chroma synthesis `syn_uv`: inputs the luma latent `[160, 16, 16]` and the chroma
/// latent `[96, 16, 16]`; output the denormalised planes U and V `[2, 256, 256]`.
pub fn syn_uv(net: &mut PtqNet) -> Res<Graph> {
let mut g = G::new();
let ly = g.b.input(checked(CHANNELS[0]), SIDE, SIDE);
let luv = g.b.input(checked(CHANNELS[1]), SIDE, SIDE);
let cat = g.concat(vec![ly, luv]);
let c1 = g.conv(net, "first_stage.conv1", cat, C3, None)?;
let x = g.concat(vec![c1, luv]);
let x = g.convt(net, "conv2_t", x, S2, 1)?;
let x = g.cab(net, x)?;
let x = g.gate(net, "iact2", x)?;
let x = g.conv(net, "conv3_t", x, P1, None)?;
let x = g.shuffle(x)?;
let x = g.tab(net, "TAM.TABs.0", x)?;
let x = g.tab(net, "TAM.TABs.1", x)?;
let x = g.gate(net, "iact3", x)?;
let x = g.convt(net, "conv4_t", x, S2, 1)?;
if g.ch(x) != 8 {
return refuse();
}
let x = g.shuffle(x)?;
let x = g.denorm(x);
g.b.finish(vec![x])
}
/// The four graphs, `rec_y`, `rec_uv`, `syn_y`, `syn_uv`, from the networks as parsed
/// (their lists move into the graphs).
pub fn nets_graphs(mut n: JpegaiNets) -> Res<[Graph; 4]> {
let [hy, huv] = &mut n.hyp;
let ry = rec(hy, Some(&mut n.mcm), false)?;
let ruv = rec(huv, None, true)?;
let [sy, suv] = &mut n.synth;
Ok([ry, ruv, syn_y(sy)?, syn_uv(suv)?])
}
/// The roles §13.0 gives these graphs, in order.
pub const NET_ROLES: [&str; 4] = ["rec_y", "rec_uv", "syn_y", "syn_uv"];
/// The four graphs by role.
pub fn nets(n: &JpegaiNets) -> Res<Lowered> {
let gs = nets_graphs(n.clone())?;
Ok(Lowered { roles: NET_ROLES.iter().copied().zip(gs).collect() })
}
/// The four graphs loaded under `p`: what [`crate::graph::load`] of each graph's bytes
/// gives, refusal for refusal. The networks' lists move into the graphs.
pub fn nets_loaded(n: JpegaiNets, p: Profile) -> Res<[Loaded; 4]> {
let [a, b, c, d] = nets_graphs(n)?;
Ok([crate::graph::load_lowered(a, p)?, crate::graph::load_lowered(b, p)?, crate::graph::load_lowered(c, p)?, crate::graph::load_lowered(d, p)?])
}

12
wai/wai-ops/src/lower/mod.rs
View file

@ -11,10 +11,11 @@ use crate::code::{Code, Refusal, Site};
use crate::graph::{CAP, Graph, InputRec, Op, Params, ScalarRec, TensorRec, op_name};
pub mod jpegai;
pub mod jpegai_nets;
mod synth;
mod whm;
pub use jpegai::{jpegai, jpegai_net, jpegai_net_loaded, parse_jpegai};
pub use jpegai::{JpegaiBytes, jpegai, jpegai_net, jpegai_net_loaded, jpegai_scale, parse_jpegai};
pub use synth::int_synth;
pub use whm::{Net, NetLayer, int_hyper, parse_wam1, parse_whm1, splat, wam1, wam1_net, wam1_net_loaded, whm1, whm1_net, whm1_net_loaded, wmb1, world, world_net, world_net_loaded};
#[cfg(kani)]
@ -69,6 +70,15 @@ impl Builder {
self.ops.push(Op { code, ins, outs: vec![t], p });
t
}
/// Add an operator writing one tensor of each of `outs`, which take the next unused
/// indices in their order; returns those tensors.
pub fn ops(&mut self, code: u8, ins: Vec<u32>, outs: Vec<TensorRec>, p: Params) -> Vec<u32> {
let first = self.tensors.len() as u32;
let ts: Vec<u32> = (first..first + outs.len() as u32).collect();
self.tensors.extend(outs);
self.ops.push(Op { code, ins, outs: ts.clone(), p });
ts
}
pub fn channels(&self, t: u32) -> u32 {
self.tensors[t as usize].channels
}

10
wai/wai-ops/tests/goldens.rs
View file

@ -306,9 +306,9 @@ fn jpegai_parsed_lowering_is_the_lowering_of_its_bytes() {
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
sets.push((rd("norm.json"), rd("hs_y.json"), rd("hs_uv.json")));
}
let committed = ["hs_y", "hs_uv"].map(|r| read(&format!("ops-conformance/certificates/wai.image.jpegai/60020d0657c4561db336a42f2d359b254a3f12348e8122aeb7bbb91d63b180ad/{r}.whg1")));
let committed = ["hs_y", "hs_uv"].map(|r| read(&format!("ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/{r}.whg1")));
for (i, (n, y, uv)) in sets.iter().enumerate() {
let bytes = lower::jpegai(n, y, uv).unwrap();
let bytes = lower::jpegai_scale(n, y, uv).unwrap();
if i == 0 {
assert_eq!(bytes.get("hs_y").unwrap().to_bytes(), committed[0], "the committed hs_y");
assert_eq!(bytes.get("hs_uv").unwrap().to_bytes(), committed[1], "the committed hs_uv");
@ -334,7 +334,7 @@ fn jpegai_parsed_lowering_is_the_lowering_of_its_bytes() {
};
for (norm, hs_y) in [(nt.replacen("\"bounds\":[", "\"bounds\":[256,", 1), yt.clone()), (nt.clone(), shift(&yt))] {
let (norm, hs_y) = (norm.into_bytes(), hs_y.into_bytes());
let a = lower::jpegai(&norm, &hs_y, uv).unwrap_err();
let a = lower::jpegai_scale(&norm, &hs_y, uv).unwrap_err();
let net = lower::parse_jpegai(&norm, &hs_y, uv).unwrap();
assert_eq!(lower::jpegai_net(&net).unwrap_err(), a);
assert_eq!(lower::jpegai_net_loaded(net, Profile::Acc32).err().unwrap(), a);
@ -352,7 +352,7 @@ fn jpegai_parsed_lowering_is_the_lowering_of_its_bytes() {
fn jpegai_scale_path_on_a_supplied_bundle() {
let dir = std::env::var("WAI_OPS_JPEGAI_BUNDLE").expect("WAI_OPS_JPEGAI_BUNDLE names the bundle");
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
let set = lower::jpegai(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
let set = lower::jpegai_scale(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
for (role, file, cz) in [("hs_y", "hs_y.json", 160u32), ("hs_uv", "hs_uv.json", 96)] {
let g = set.get(role).unwrap();
let net = json::parse(&rd(file)).unwrap();
@ -428,7 +428,7 @@ fn jpegai_entropy_path_matches_the_legacy_decoder() {
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
let norm = json::parse(&rd("norm.json")).unwrap();
let d = json::parse(&std::fs::read(&dump).unwrap()).unwrap();
let set = lower::jpegai(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
let set = lower::jpegai_scale(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
let pmf: Vec<Vec<i64>> = norm.get("pmf").unwrap().as_arr().unwrap().iter().map(i64s).collect();
let tables = stage::mans::build_latent_tables(&pmf).unwrap();
let want: Vec<Vec<i64>> = d.get("tables").unwrap().as_arr().unwrap().iter().map(i64s).collect();

19
wai/wai-rs/Cargo.toml
View file

@ -819,7 +819,7 @@ required-features = ["neural_int", "video_int", "audio_int", "world_int", "splat
[[test]]
name = "int_kernels_geometry"
path = "tests/int_kernels_geometry.rs"
required-features = ["neural_int"]
required-features = ["neural_int", "legacy_int"]
[[test]]
name = "int_ops_differential"
@ -857,6 +857,11 @@ name = "jpegai_ops"
path = "tests/jpegai_ops.rs"
required-features = ["neural_int", "legacy_int"]
[[test]]
name = "jpegai_nets_ops"
path = "tests/jpegai_nets_ops.rs"
required-features = ["neural_int", "legacy_int"]
[[test]]
name = "int_video_ops"
path = "tests/int_video_ops.rs"
@ -1165,7 +1170,7 @@ required-features = ["neural_int", "legacy_int"]
[[test]]
name = "jpegai_e2e_synthesis"
path = "tests/jpegai_e2e_synthesis.rs"
required-features = ["neural_int"]
required-features = ["neural_int", "legacy_int"]
[[test]]
name = "jpegai_efl_calibration"
@ -1180,7 +1185,7 @@ required-features = ["neural_int"]
[[test]]
name = "jpegai_hyper_calibration"
path = "tests/jpegai_hyper_calibration.rs"
required-features = ["neural"]
required-features = ["neural", "legacy_int"]
[[test]]
name = "jpegai_hyperscale_calibration"
@ -1225,7 +1230,7 @@ required-features = ["neural"]
[[test]]
name = "jpegai_mcm_calibration"
path = "tests/jpegai_mcm_calibration.rs"
required-features = ["neural"]
required-features = ["neural", "legacy_int"]
[[test]]
name = "jpegai_postfilter_chain"
@ -1235,7 +1240,7 @@ required-features = ["neural_int"]
[[test]]
name = "jpegai_reconstruct_calibration"
path = "tests/jpegai_reconstruct_calibration.rs"
required-features = ["neural_int"]
required-features = ["neural_int", "legacy_int"]
[[test]]
name = "jpegai_reconstruction_binding"
@ -1255,12 +1260,12 @@ required-features = ["neural_int", "legacy_int"]
[[test]]
name = "jpegai_synthesis_calibration"
path = "tests/jpegai_synthesis_calibration.rs"
required-features = ["neural"]
required-features = ["neural", "legacy_int"]
[[test]]
name = "jpegai_uv_calibration"
path = "tests/jpegai_uv_calibration.rs"
required-features = ["neural"]
required-features = ["neural", "legacy_int"]
[[test]]
name = "neural_cost_limit"

62
wai/wai-rs/src/bin/wai_decode_bench.rs
View file

@ -206,7 +206,7 @@ fn jpegai_weights(luma: bool) -> wai::jpegai_synthesis::Weights {
} else {
conv(&mut m, "conv4_t", [64, 8, 3, 3], 64 * 9 / 4);
}
wai::jpegai_synthesis::Weights::from_json(&Value::Object(m))
wai::jpegai_synthesis::Weights::from_json(&Value::Object(m)).unwrap()
}
/// `model` lowered by `whm1/1` (with the RGB8 step where `rgb8` is set) and loaded under
@ -224,8 +224,32 @@ fn run_lowered(g: &wai_ops::Loaded, x: &[i64], c: usize, h: usize, w: usize) ->
// ─────────────────────────────── cases ───────────────────────────────
/// One convolution (or transposed convolution) as a `WHG1` graph with one operator, its
/// input checked within i32 at `h × h` and its `RQ` the JPEG-AI networks' (one
/// multiplier, one shift, a `post` bias, `floor`; integer-operators §13.6.2), loaded
/// under `acc64`: the operator the decode runs.
#[allow(clippy::too_many_arguments)]
fn conv_graph(cin: usize, cout: usize, k: usize, stride: usize, pad: usize, groups: usize, opad: usize, transposed: bool, h: usize, w: Vec<i64>, mult: i64, bias: Vec<i64>) -> wai_ops::Loaded {
use wai_ops::graph::{opcode, BiasKind, Graph, InputRec, Op, Params, Round, Rq, TensorRec, CAP};
let u = |v: usize| v as u32;
let rq = Rq { round: Round::Floor, bias_kind: BiasKind::Post, m_kind: 1, s_kind: 1, m_bytes: 8, b_bytes: 8, mult: vec![mult], shift: vec![31], bias };
let (code, p) = if transposed {
(opcode::CONV_TRANSPOSE, Params::ConvT { cout: u(cout), kh: u(k), kw: u(k), sh: u(stride), sw: u(stride), ph: u(pad), pw: u(pad), oph: u(opad), opw: u(opad), w_bytes: 2, rq, w })
} else {
(opcode::CONV, Params::Conv { cout: u(cout), kh: u(k), kw: u(k), sh: u(stride), sw: u(stride), ph: u(pad), pw: u(pad), groups: u(groups), w_bytes: 2, rq, w })
};
let g = Graph {
max_elems: CAP as u32,
scalars: vec![],
inputs: vec![InputRec { tensor: 0, max_h: u(h), max_w: u(h), kind: 0, src: 0, init: 0 }],
outputs: vec![1],
tensors: vec![TensorRec::check(u(cin), i32::MIN as i64, i32::MAX as i64), TensorRec::derived(u(cout))],
ops: vec![Op { code, ins: vec![0], outs: vec![1], p }],
};
wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64).expect("the convolution loads under acc64")
}
fn kernel_cases(cases: &mut Vec<Case>) {
use wai::int_transformer::{conv_grouped, conv_transpose};
// int_hyper g_s middle layer: convT 192→192, 5×5, stride 2, 64² → 128².
{
@ -263,7 +287,9 @@ fn kernel_cases(cases: &mut Vec<Case>) {
run: Box::new(move || Out::I(wai::int_transform::igdn(&x, 192, 128 * 128, 12, &gamma, &beta, 16))),
});
}
// JPEG-AI shapes (Q16 activations, 16-bit weights, shift 31).
// JPEG-AI shapes (Q16 activations, 16-bit weights, shift 31), each convolution as the
// one-operator graph the decode runs (through the interpreter with the fast kernels
// behind it).
type Spec = (&'static str, usize, usize, usize, usize, usize, usize, usize, bool);
// name, cin, cout, k, stride, groups, h(=w) in, transposed?
let specs: [Spec; 5] = [
@ -278,23 +304,17 @@ fn kernel_cases(cases: &mut Vec<Case>) {
let cin_g = cin / groups;
let x = g.vec(cin * hw * hw, 2 << 16);
let wt = g.vec(cin_g * cout * k * k, 32767);
let mult = vec![(((3.0 / (cin_g * k * k) as f64).sqrt() / 32767.0) * (1u64 << 31) as f64) as i64; cout];
let mult = (((3.0 / (cin_g * k * k) as f64).sqrt() / 32767.0) * (1u64 << 31) as f64) as i64;
let bias = g.vec(cout, 1 << 12);
let pad = k / 2;
let graph = conv_graph(cin, cout, k, stride, pad, groups, opad, transposed, hw, wt, mult, bias);
let ho = if transposed { (hw - 1) * stride + k + opad - 2 * pad } else { (hw + 2 * pad - k) / stride + 1 };
let macs = if transposed { cin * cout * k * k * hw * hw } else { cin_g * cout * k * k * ho * ho } as u64;
cases.push(Case {
name: format!("kernel · JPEG-AI {label}"),
per_run: ((ho * ho) as f64 / 1e6, "MPix"),
macs: Some(macs),
run: Box::new(move || {
let out = if transposed {
conv_transpose(&x, cin, hw, hw, &wt, cout, k, stride, pad, opad, &mult, &bias, 31).0
} else {
conv_grouped(&x, cin, hw, hw, &wt, cout, k, stride, pad, groups, &mult, &bias, 31).0
};
Out::I(out)
}),
run: Box::new(move || Out::I(run_lowered(&graph, &x, cin, hw, hw).data)),
});
}
}
@ -314,9 +334,15 @@ fn decoder_cases(cases: &mut Vec<Case>, quick: bool) {
run: Box::new(move || Out::B(run_lowered(&model, &y, 320, lat, lat).data.iter().map(|&v| v as u8).collect())),
});
}
// JPEG-AI synthesis (luma + chroma nets), 256² (the supported geometry).
// JPEG-AI synthesis (luma + chroma nets), 256² (the supported geometry), through the
// graphs the decode runs (`syn_y`, `syn_uv`; integer-operators §13.6.2).
if !quick {
let (wy, wuv) = (jpegai_weights(true), jpegai_weights(false));
let lower = |luma: bool| -> wai_ops::Loaded {
let mut net = jpegai_weights(luma).ptq().expect("the synthetic weights convert");
let g = if luma { wai_ops::lower::jpegai_nets::syn_y(&mut net) } else { wai_ops::lower::jpegai_nets::syn_uv(&mut net) };
wai_ops::graph::load(&g.expect("the synthetic networks lower").to_bytes(), wai_ops::Profile::Acc64).expect("they load under acc64")
};
let (gy, guv) = (lower(true), lower(false));
let mut g = Rng(0x1A_0003);
let latent = g.vec(256 * 16 * 16, 1 << 16); // y‖uv latent, Q16
cases.push(Case {
@ -324,9 +350,11 @@ fn decoder_cases(cases: &mut Vec<Case>, quick: bool) {
per_run: (256.0 * 256.0 / 1e6, "MPix"),
macs: None,
run: Box::new(move || {
let mut y = wai::jpegai_synthesis::decode_y(&latent[..160 * 256], &wy);
y.extend(wai::jpegai_synthesis::decode_uv(&latent, &wuv));
Out::B(y)
let t = |c: u32, d: &[i64]| wai_ops::Tensor::new(c, 16, 16, d.to_vec());
let run = |g: &wai_ops::Loaded, x: Vec<wai_ops::Tensor>| wai_ops::run_with(g, x, &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host).expect("the latent is within its ranges").swap_remove(0).data;
let mut y = run(&gy, vec![t(160, &latent[..160 * 256])]);
y.extend(run(&guv, vec![t(160, &latent[..160 * 256]), t(96, &latent[160 * 256..])]));
Out::I(y)
}),
});
}

31
wai/wai-rs/src/codecs/mod.rs
View file

@ -1348,24 +1348,28 @@ pub struct OpsMember {
pub rans64: Option<(&'static str, &'static str, &'static str)>,
/// `bucket`'s rule (E3), where the driver buckets a scale.
pub bucket: Option<&'static str>,
/// Roles of the lowering whose graphs every conforming set holds to a stronger profile
/// than `profile`, each with that profile (§11, `role_profiles`); empty for most.
pub role_profiles: &'static [(&'static str, &'static str)],
}
/// The `ops` member of `cap`, for a capability decoded through the integer operator
/// set; `None` for any other. A capability gains one when its decoder runs its
/// networks through their lowering (`extensions/integer-operators.md` §15.1), or, with
/// no network, its stage operators (§13.5).
/// `wai.image.jpegai` has one for its scale path and entropy decodes (§13.6.1), the
/// stages its `entropy-consistency` claim covers.
/// `wai.image.jpegai` has one for its integer stages (§13.6): its scale path and entropy
/// decodes, the stages its `entropy-consistency` claim covers, and its other integer
/// networks.
/// `wai.video.int_motion` and the four `int_refine` capabilities run stage operators
/// only: `WIV1` keyframe 0 E1, E10 and E9; `WIR1` E1, E11, E12 and E9.
/// `wai.world.model` has none: its integer renderer runs through `world/1`, which is
/// informative (§13.5), and the capability's decode is the replay of its action trace.
pub fn capability_ops(cap: &str) -> Option<OpsMember> {
let ops = |lowering, rans64, bucket| Some(OpsMember { profile: Some("acc64"), subset: Some("full"), lowering: Some(lowering), rans64, bucket });
let ops = |lowering, rans64, bucket| Some(OpsMember { profile: Some("acc64"), subset: Some("full"), lowering: Some(lowering), rans64, bucket, role_profiles: &[] });
// integer-payloads §2's stream: `rans64_decode` with `t4097`, `lenient` and `canonical`.
const PAYLOAD_STREAM: Option<(&str, &str, &str)> = Some(("t4097", "lenient", "canonical"));
// A decode with no graph: its stage operators, and the modes of its E1.
let stages = |rans64| Some(OpsMember { profile: None, subset: None, lowering: None, rans64, bucket: None });
let stages = |rans64| Some(OpsMember { profile: None, subset: None, lowering: None, rans64, bucket: None, role_profiles: &[] });
match cap {
CAP_NEURAL_INT_SYNTH => ops("int_synth/1", None, None),
CAP_SPLAT_INT_CODEC => ops("whm1/1", PAYLOAD_STREAM, None),
@ -1378,11 +1382,20 @@ pub fn capability_ops(cap: &str) -> Option<OpsMember> {
// `final_only` end rule, and each latent's σ is bucketed with E3's `ceil` rule
// (integer-payloads §14).
CAP_NEURAL_INT_MLICPP => ops("wmb1/1", Some(("t2", "lenient", "final_only")), Some("ceil")),
// JPEG AI's scale path and its two entropy decodes (integer-operators §13.6.1): the
// scale graphs of every conforming bundle meet `acc32`, so the profile is `acc32`;
// its stage operators (`mans_decode`, `coded`, `block_stat`, `band_lookup`,
// `scale_index`) have no modes to fix.
CAP_IMAGE_JPEGAI => Some(OpsMember { profile: Some("acc32"), subset: Some("full"), lowering: Some("jpegai-int/1"), rans64: None, bucket: None }),
// JPEG AI's integer stages (integer-operators §13.6): every graph of a conforming
// bundle meets `acc64`, and its scale graphs `hs_y` and `hs_uv` meet `acc32`
// (§13.6.1), so they run on a 32-bit fabric as on a 64-bit one; its other
// networks (§13.6.2) use operators that never meet `acc32` (§7.3). Its stage
// operators (`mans_decode`, `coded`, `block_stat`, `band_lookup`, `scale_index`)
// have no modes to fix.
CAP_IMAGE_JPEGAI => Some(OpsMember {
profile: Some("acc64"),
subset: Some("full"),
lowering: Some("jpegai-int/1"),
rans64: None,
bucket: None,
role_profiles: &[("hs_y", "acc32"), ("hs_uv", "acc32")],
}),
// No graph (§13.5): a `WIV1` clip's residual stream and a `WIR1` unit's stream are
// §2's, and the predictors and the step (E10, E11, E12, E9) have no modes.
CAP_VIDEO_INT_MOTION | CAP_IMAGE_INT_REFINE | CAP_VIDEO_INT_REFINE | CAP_AUDIO_INT_REFINE | CAP_SPLAT_INT_REFINE => stages(PAYLOAD_STREAM),

25
wai/wai-rs/src/codecs/registry.rs
View file

@ -853,6 +853,9 @@ fn ops_json(o: &super::OpsMember) -> serde_json::Value {
v[k] = Value::from(m);
}
}
if !o.role_profiles.is_empty() {
v["role_profiles"] = Value::Object(o.role_profiles.iter().map(|(r, p)| (r.to_string(), Value::from(*p))).collect());
}
let mut stage = Map::new();
if let Some((table, escape, end)) = o.rans64 {
stage.insert("rans64".into(), json!({ "table": table, "escape": escape, "end": end }));
@ -882,7 +885,8 @@ fn ops_schema() -> serde_json::Value {
"dependentRequired": {
"profile": ["subset", "lowering"],
"subset": ["profile", "lowering"],
"lowering": ["profile", "subset"]
"lowering": ["profile", "subset"],
"role_profiles": ["profile"]
},
"anyOf": [{ "required": ["lowering"] }, { "required": ["stage"] }],
"properties": {
@ -890,6 +894,12 @@ fn ops_schema() -> serde_json::Value {
"profile": { "enum": ["acc32", "acc64"] },
"subset": { "enum": ["full", "npu"] },
"lowering": { "type": "string" },
"role_profiles": {
"type": "object",
"description": "Roles of the lowering whose graphs every conforming set holds to a stronger profile than `profile`, each with that profile (extensions/integer-operators.md §11).",
"minProperties": 1,
"additionalProperties": { "enum": ["acc32", "acc64"] }
},
"stage": {
"type": "object",
"additionalProperties": false,
@ -1511,10 +1521,12 @@ mod tests {
assert_eq!(o.rans64, Some(("t4097", "lenient", "canonical")), "{}", r.capability);
continue;
}
// JPEG AI's scale path and entropy decodes, which its entropy-consistency claim
// covers, decode through the operator set at acc32 (§13.6.1).
let (class, profile) = if r.capability == CAP_IMAGE_JPEGAI { (ReplicateClass::NeuralEntropyExact, "acc32") } else { (ReplicateClass::NeuralIntegerExact, "acc64") };
// JPEG AI's integer stages decode through the operator set at acc64, its scale
// path, which its entropy-consistency claim covers, at acc32 (§13.6).
let (class, profile) = if r.capability == CAP_IMAGE_JPEGAI { (ReplicateClass::NeuralEntropyExact, "acc64") } else { (ReplicateClass::NeuralIntegerExact, "acc64") };
assert_eq!(r.replicate, Some(class), "{}", r.capability);
let roles: &[(&str, &str)] = if r.capability == CAP_IMAGE_JPEGAI { &[("hs_y", "acc32"), ("hs_uv", "acc32")] } else { &[] };
assert_eq!(o.role_profiles, roles, "{}", r.capability);
assert_eq!((o.profile, o.subset), (Some(profile), Some("full")), "{}", r.capability);
assert!(o.lowering.is_some_and(|l| ["int_synth/1", "whm1/1", "wam1/1", "wmb1/1", "world/1", "jpegai-int/1"].contains(&l)), "{}", r.capability);
}
@ -1538,7 +1550,10 @@ mod tests {
"set": "wai-ops/1",
"stage": { "rans64": { "table": "t4097", "escape": "lenient", "end": "canonical" } }
}),
CAP_IMAGE_JPEGAI => serde_json::json!({ "set": "wai-ops/1", "profile": "acc32", "subset": "full", "lowering": "jpegai-int/1" }),
CAP_IMAGE_JPEGAI => serde_json::json!({
"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "jpegai-int/1",
"role_profiles": { "hs_y": "acc32", "hs_uv": "acc32" }
}),
CAP_NEURAL_INT_MLICPP => serde_json::json!({
"set": "wai-ops/1", "profile": "acc64", "subset": "full", "lowering": "wmb1/1",
"stage": { "rans64": { "table": "t2", "escape": "lenient", "end": "final_only" }, "bucket": "ceil" }

11
wai/wai-rs/src/int_kernels.rs
View file

@ -302,12 +302,16 @@ pub(crate) fn par_for_each<T: Send, F: Fn(T) + Sync>(items: Vec<T>, threads: usi
}
}
/// Elements per work item of the element-wise helpers.
/// Elements per work item of the element-wise helpers. The helpers below serve the
/// direct JPEG AI networks (`jpegai_synthesis`), which are test-only (`legacy_int`): the
/// decode runs those networks through the integer operator set.
#[cfg(any(test, feature = "legacy_int"))]
const MAP_CHUNK: usize = 1 << 14;
/// Element-wise `out[i] = f(x[i])`, chunked over workers when large. Every element
/// is computed by the same pure function of its own input, so the partition cannot
/// change any output. Serial under the reference policy (one thread).
#[cfg(any(test, feature = "legacy_int"))]
pub(crate) fn par_map<F: Fn(i64) -> i64 + Sync>(x: &[i64], f: F) -> Vec<i64> {
let mut out = vec![0i64; x.len()];
let threads = threads_for(x.len() as u64 * 16);
@ -321,6 +325,7 @@ pub(crate) fn par_map<F: Fn(i64) -> i64 + Sync>(x: &[i64], f: F) -> Vec<i64> {
/// Element-wise `out[i] = f(a[i], b[i])` over `min(a.len(), b.len())` elements (the
/// length `zip` gives), chunked like [`par_map`].
#[cfg(any(test, feature = "legacy_int"))]
pub(crate) fn par_zip<F: Fn(i64, i64) -> i64 + Sync>(a: &[i64], b: &[i64], f: F) -> Vec<i64> {
let n = a.len().min(b.len());
let mut out = vec![0i64; n];
@ -1138,7 +1143,9 @@ fn note_fast_conv() {
/// overflow there panics the reference in an overflow-checked build and would pass
/// unnoticed here. Under the bound the fast kernel uses `i32×i32→i64` products (and
/// SIMD lanes) when every operand fits in `i32`, else the reference's own `i64`
/// arithmetic; either way each output element is the same integer.
/// arithmetic; either way each output element is the same integer. Test-only
/// (`legacy_int`), with the direct networks that call it.
#[cfg(any(test, feature = "legacy_int"))]
pub(crate) fn conv_acc_i64<E>(g: &ConvGeom, x: &[i64], wt: &[i64], epi: E) -> Option<Vec<i64>>
where
E: Fn(usize, i64) -> i64 + Sync,

7
wai/wai-rs/src/int_set_cache.rs
View file

@ -83,6 +83,13 @@ impl Resident for crate::jpegai_ops::ScalePath {
}
}
/// `wai.image.jpegai`'s other integer networks (`jpegai_ops::Networks`): its four graphs.
impl Resident for crate::jpegai_ops::Networks {
fn heap_bytes(&self) -> u64 {
crate::jpegai_ops::Networks::heap_bytes(self)
}
}
/// The key of a loaded set: the capability whose driver loaded it, the lowering
/// (integer-operators §13), the profile its graphs were loaded under (§11), and the
/// SHA-256 of the members it was lowered from, as [`SetKey::of`] forms it.

147
wai/wai-rs/src/jpegai_decode.rs
View file

@ -60,11 +60,10 @@ use crate::jpegai_bitstream::{find_segment, parse_container, parse_lsbs_flags,
parse_picture_header, parse_tools_header, Marker, PictureHeader,
ToolsHeader};
use crate::jpegai_icci::IcciNet;
use crate::jpegai_ops::{Plane, ScalePath};
use crate::jpegai_ops::{Networks, Plane, ScalePath};
use crate::jpegai_pipeline::{dequantize_resi, efl_chroma, efn_filter, gain_unit_scaler_vec,
lef_filter, lsbs_refine, reconstruct_yhat, reconstruct_yhat_uv};
use crate::jpegai_synthesis::{decode_hyper, decode_uv_f32, decode_y_f32,
q16_from_symbols, IntNet, Weights};
lef_filter, lsbs_refine};
use crate::jpegai_synthesis::{IntNet, Weights};
/// One plane's model weights: the fully-integer hyper-scale net, the PTQ hyper decoder,
/// and the PTQ synthesis net.
@ -117,23 +116,37 @@ fn q16(f: &[f32]) -> Vec<i64> { f.iter().map(|&x| round_half_away(x as f64 * 655
/// Where a plane's integers (integer-operators §13.6.1 steps 1–4) come from.
#[derive(Clone, Copy)]
enum Front<'a> {
/// The direct computation ([`crate::jpegai_ops::legacy`], test-only).
/// The direct computation ([`crate::jpegai_ops::legacy`], test-only), from the
/// bundle as parsed.
#[cfg(any(test, feature = "legacy_int"))]
Direct,
Direct(&'a DecoderWeights),
/// The integer operator set, from the bundle's scale path as loaded.
Ops(&'a ScalePath),
}
/// Where a plane's other integer networks (integer-operators §13.6.2) run: the hyper
/// decoder and the context stages, and the synthesis.
#[derive(Clone, Copy)]
enum Back<'a> {
/// The direct networks ([`crate::jpegai_synthesis`], [`crate::jpegai_pipeline`],
/// test-only), from the bundle as parsed.
#[cfg(any(test, feature = "legacy_int"))]
Direct(&'a DecoderWeights),
/// The integer operator set, from the bundle's networks as loaded.
Ops(&'a Networks),
}
/// A refusal of the exact decode.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Refusal {
/// The codestream, the target or the floating-point environment is outside what the
/// decoder implements ([`check_scope`]'s reason).
Scope(&'static str),
/// A refusal of the integer stages (integer-operators §13.6.1): `model_mismatch` for
/// a bundle that fails a load check or whose scale graphs do not meet the
/// capability's profile, `stream_exhausted` for an entropy substream the decode
/// cannot read (E8).
/// A refusal of the integer stages (integer-operators §13.6): `model_mismatch` for a
/// bundle that fails a load check, or whose graphs do not meet their profiles (`acc32`
/// for the scale path, `acc64` for the networks); `stream_exhausted` for an entropy
/// substream the decode cannot read (E8); `field_range` for a dequantised residual, a
/// latent or an activation of the networks that leaves i32 (§13.6.2).
Ops(wai_ops::Code),
}
@ -152,8 +165,9 @@ impl std::fmt::Display for Refusal {
match self {
Refusal::Scope(why) => write!(f, "unsupported codestream: {why}"),
Refusal::Ops(c) => match c {
wai_ops::Code::ModelMismatch => write!(f, "model_mismatch: the bundle fails a load check of integer-operators §13.6.1, or its scale graphs do not meet acc32"),
wai_ops::Code::ModelMismatch => write!(f, "model_mismatch: the bundle fails a load check of integer-operators §13.6, or its graphs do not meet their profiles (acc32 for the scale path, acc64 for the networks)"),
wai_ops::Code::StreamExhausted => write!(f, "stream_exhausted: an entropy substream ends before the decode it carries"),
wai_ops::Code::FieldRange => write!(f, "field_range: a dequantised residual, a latent or an activation of the integer networks leaves i32 (integer-operators §13.6.2)"),
other => write!(f, "{}", other.as_str()),
},
}
@ -178,8 +192,7 @@ fn payload<'a>(data: &'a [u8], segs: &[crate::jpegai_bitstream::Segment], marker
/// LEF's sigma surface.
#[allow(clippy::too_many_arguments)]
fn decode_plane(
data: &[u8], plane: Plane, n: &NormTables, pw: &PlaneWeights,
gain: &[i64], mcm: Option<[&Weights; 4]>, front: Front<'_>,
data: &[u8], plane: Plane, n: &NormTables, gain: &[i64], front: Front<'_>, back: Back<'_>,
) -> Result<(Vec<f32>, Vec<i64>, Vec<i64>), Refusal> {
const H: usize = 16; // the latent plane is 16×16; z is at 4×4 and the nets upsample ×4
const W: usize = 16;
@ -193,7 +206,7 @@ fn decode_plane(
// SOZ → z → v → mask, likely, rvs, σ → SORP/SORS → y_q
let ints = match front {
#[cfg(any(test, feature = "legacy_int"))]
Front::Direct => crate::jpegai_ops::legacy::plane(n, &pw.hyper_scale, gain, plane, soz, latent, &grfs),
Front::Direct(w) => crate::jpegai_ops::legacy::plane(n, &plane_weights(w, plane).hyper_scale, gain, plane, soz, latent, &grfs),
Front::Ops(sp) => crate::jpegai_ops::plane(sp, plane, soz, latent, &grfs)?,
};
// Every symbol is a 16-bit value (E8's wrap16).
@ -202,17 +215,34 @@ fn decode_plane(
// dequant → resi ; ψ ; reconstruct → ŷ ; LSBS → latent
let sv = gain_unit_scaler_vec(gain, 0);
let resi = dequantize_resi(&y_q, &ints.likely, &grfs, &sv, channels);
let (psi, ..) = decode_hyper(&q16_from_symbols(&ints.z), channels, 4, 4, &pw.hyper_decoder);
let resi_q = q16(&resi);
let yhat = match mcm {
Some(m) => reconstruct_yhat(&resi_q, &psi, channels, H, W, m),
None => reconstruct_yhat_uv(&resi_q, &psi, channels, H, W),
let yhat = match back {
#[cfg(any(test, feature = "legacy_int"))]
Back::Direct(w) => {
use crate::jpegai_pipeline::{reconstruct_yhat, reconstruct_yhat_uv};
use crate::jpegai_synthesis::{decode_hyper, q16_from_symbols};
let (psi, ..) = decode_hyper(&q16_from_symbols(&ints.z), channels, 4, 4, &plane_weights(w, plane).hyper_decoder);
match plane {
Plane::Luma => reconstruct_yhat(&resi_q, &psi, channels, H, W, [&w.mcm_y[0], &w.mcm_y[1], &w.mcm_y[2], &w.mcm_y[3]]),
Plane::Chroma => reconstruct_yhat_uv(&resi_q, &psi, channels, H, W),
}
}
Back::Ops(nets) => crate::jpegai_ops::reconstruct(nets, plane, ints.z.clone(), resi_q)?,
};
let yhat_f: Vec<f32> = yhat.iter().map(|&v| v as f64 as f32 / 65536.0).collect();
let lat = lsbs_refine(&yhat_f, &resi, &ints.likely, &n.lsbs_t0, &n.lsbs_t1, channels, H, W);
Ok((lat, ints.v, ints.rvs))
}
/// A plane's weights in a parsed bundle.
#[cfg(any(test, feature = "legacy_int"))]
fn plane_weights(w: &DecoderWeights, p: Plane) -> &PlaneWeights {
match p {
Plane::Luma => &w.luma,
Plane::Chroma => &w.chroma,
}
}
fn grfs_flags(data: &[u8], channels: usize, marker: Marker) -> Vec<bool> {
let segs = parse_container(data);
let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
@ -369,31 +399,42 @@ pub fn check_scope(data: &[u8]) -> Result<(), &'static str> {
scope(data).map(|_| ())
}
/// Decode a JPEG-AI codestream to RGB. The scale path and the two entropy decodes run
/// through the integer operator set (`extensions/integer-operators.md` §13.6.1, lowered
/// by `jpegai-int/1`): the bundle's scale graphs are lowered and loaded under `acc32`
/// with its tables ([`load`]), then each plane's integers come from `wai-ops`.
/// Decode a JPEG-AI codestream to RGB. Every integer stage runs through the integer
/// operator set (`extensions/integer-operators.md` §13.6, lowered by `jpegai-int/1`):
/// the bundle's scale graphs are lowered and loaded under `acc32` with its tables
/// ([`load`]), and its hyper decoders, context stages and synthesis under `acc64`
/// ([`load_networks`]); then each plane's integers, its reconstructed latent and the
/// synthesis come from `wai-ops`.
///
/// Refuses ([`Refusal`]) a codestream [`check_scope`] refuses, with its reason; a bundle
/// that fails a load check of §13.6.1 or whose scale graphs do not meet `acc32`
/// (`model_mismatch`); and an entropy substream the decode cannot read
/// (`stream_exhausted`). It does not panic on a stream.
/// that fails a load check of §13.6 or whose graphs do not meet their profiles
/// (`model_mismatch`); an entropy substream the decode cannot read (`stream_exhausted`);
/// and a dequantised residual, a latent or an activation of the networks past i32
/// (`field_range`). It does not panic on a stream.
pub fn decode(data: &[u8], w: &DecoderWeights, n: &NormTables) -> Result<Rgb, Refusal> {
decode_with_stages(data, w, n, &mut |_, _| {})
}
/// The bundle's scale path, lowered and loaded from the bundle as parsed
/// ([`crate::jpegai_ops::load`]): a sink that decodes several codestreams against one
/// bundle loads it once and runs [`decode_loaded`].
/// bundle loads it once, with its networks ([`load_networks`]), and runs
/// [`decode_loaded`].
pub fn load(w: &DecoderWeights, n: &NormTables) -> Result<ScalePath, Refusal> {
Ok(crate::jpegai_ops::load(n, &w.luma.hyper_scale, &w.chroma.hyper_scale)?)
}
/// [`decode_with_stages`] with the bundle's scale path already loaded (`sp`, from the
/// members `w` and `n` were parsed from).
pub fn decode_loaded(data: &[u8], w: &DecoderWeights, n: &NormTables, sp: &ScalePath,
/// The bundle's other integer networks (integer-operators §13.6.2), lowered and loaded
/// from the bundle as parsed ([`crate::jpegai_ops::load_networks`], under `acc64`).
pub fn load_networks(w: &DecoderWeights) -> Result<Networks, Refusal> {
Ok(crate::jpegai_ops::load_networks(w)?)
}
/// [`decode_with_stages`] with the bundle's scale path and networks already loaded (`sp`
/// and `nets`, from the members `n` and the bundle were parsed from). Of the bundle's
/// float stages it reads the eICCI net `icci`; every integer stage is `sp`'s and `nets`'.
pub fn decode_loaded(data: &[u8], icci: &IcciNet, n: &NormTables, sp: &ScalePath, nets: &Networks,
observe: &mut dyn FnMut(Stage, &[&[f32]])) -> Result<Rgb, Refusal> {
run(data, w, n, Front::Ops(sp), observe)
run(data, icci, n, Front::Ops(sp), Back::Ops(nets), observe)
}
/// A point in the decode at which [`decode_with_stages`] reports the float values the
@ -438,18 +479,21 @@ impl Stage {
/// (a final byte hides float differences smaller than its rounding step). `observe`
/// sees the planes in the order Y, U, V (one slice for a latent stage). It does not
/// change the decode: the returned image is the one [`decode`] returns. The codestream's
/// scope is checked first, then the bundle's scale path is loaded, then the decode runs.
/// scope is checked first, then the bundle's scale path is loaded, then its networks,
/// then the decode runs.
pub fn decode_with_stages(data: &[u8], w: &DecoderWeights, n: &NormTables,
observe: &mut dyn FnMut(Stage, &[&[f32]])) -> Result<Rgb, Refusal> {
scope(data).map_err(Refusal::Scope)?;
let sp = load(w, n)?;
decode_loaded(data, w, n, &sp, observe)
let nets = load_networks(w)?;
decode_loaded(data, &w.icci, n, &sp, &nets, observe)
}
/// The decode as it ran before its integer stages ran through the operator set
/// (test-only, `legacy_int`): each plane's integers from the direct computation
/// ([`crate::jpegai_ops::legacy`]), which panics on a stream that runs out. `None` where
/// the decode refused.
/// ([`crate::jpegai_ops::legacy`]), which panics on a stream that runs out, and the
/// direct hyper decoders, context stages and synthesis, which compute on past the
/// operator set's ranges. `None` where the decode refused.
#[cfg(any(test, feature = "legacy_int"))]
pub mod legacy {
use super::*;
@ -462,17 +506,17 @@ pub mod legacy {
/// The direct [`super::decode_with_stages`].
pub fn decode_with_stages(data: &[u8], w: &DecoderWeights, n: &NormTables,
observe: &mut dyn FnMut(Stage, &[&[f32]])) -> Option<Rgb> {
run(data, w, n, Front::Direct, observe).ok()
run(data, &w.icci, n, Front::Direct(w), Back::Direct(w), observe).ok()
}
}
/// The decode, with each plane's integers from `front`.
fn run(data: &[u8], w: &DecoderWeights, n: &NormTables, front: Front<'_>,
fn run(data: &[u8], icci: &IcciNet, n: &NormTables, front: Front<'_>, back: Back<'_>,
observe: &mut dyn FnMut(Stage, &[&[f32]])) -> Result<Rgb, Refusal> {
let (ph, th) = scope(data).map_err(Refusal::Scope)?;
let (gain_y, gain_uv) = match front {
#[cfg(any(test, feature = "legacy_int"))]
Front::Direct => {
Front::Direct(_) => {
if n.gain_y.len() != 160 || n.gain_uv.len() != 96 || n.per_channel == 0 {
return Err(Refusal::Ops(wai_ops::Code::ModelMismatch));
}
@ -483,16 +527,29 @@ fn run(data: &[u8], w: &DecoderWeights, n: &NormTables, front: Front<'_>,
let (h, w_) = (ph.img_height as usize, ph.img_width as usize);
// ---- luma + chroma latents → pixel planes (float, pre-round) ----
let mcm = [&w.mcm_y[0], &w.mcm_y[1], &w.mcm_y[2], &w.mcm_y[3]];
let (y_lat, y_v, y_rvs) = decode_plane(data, Plane::Luma, n, &w.luma, gain_y, Some(mcm), front)?;
let (y_lat, y_v, y_rvs) = decode_plane(data, Plane::Luma, n, gain_y, front, back)?;
observe(Stage::LatentY, &[&y_lat]);
let (uv_lat, ..) = decode_plane(data, Plane::Chroma, n, &w.chroma, gain_uv, None, front)?;
let (uv_lat, ..) = decode_plane(data, Plane::Chroma, n, gain_uv, front, back)?;
observe(Stage::LatentUv, &[&uv_lat]);
let mut yp = decode_y_f32(&q16(&y_lat), &w.luma.synthesis);
let mut cat = q16(&y_lat);
cat.extend(q16(&uv_lat));
let uv = decode_uv_f32(&cat, &w.chroma.synthesis);
// The synthesis: the latents in Q16 in, the planes in Q16 out, read as `v / 2^16`.
let (mut yp, uv) = match back {
#[cfg(any(test, feature = "legacy_int"))]
Back::Direct(w) => {
use crate::jpegai_synthesis::{decode_uv_f32, decode_y_f32};
let yp = decode_y_f32(&q16(&y_lat), &w.luma.synthesis);
let mut cat = q16(&y_lat);
cat.extend(q16(&uv_lat));
(yp, decode_uv_f32(&cat, &w.chroma.synthesis))
}
Back::Ops(nets) => {
let (ly, luv) = (q16(&y_lat), q16(&uv_lat));
let yp = crate::jpegai_ops::synthesize_y(nets, ly.clone())?;
let uv = crate::jpegai_ops::synthesize_uv(nets, ly, luv)?;
let f = |v: Vec<i64>| -> Vec<f32> { v.into_iter().map(|x| x as f32 / 65536.0).collect() };
(f(yp), f(uv))
}
};
let (mut up, mut vp) = (uv[..h * w_].to_vec(), uv[h * w_..].to_vec());
observe(Stage::Synthesis, &[&yp, &up, &vp]);
@ -508,7 +565,7 @@ fn run(data: &[u8], w: &DecoderWeights, n: &NormTables, front: Front<'_>,
// eICCI — [0,1] domain: normalise in, scale back out.
if th.icci_enabled {
let nrm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::<Vec<f32>>();
let (yo, uo, vo) = w.icci.forward(&nrm(&yp), &nrm(&up), &nrm(&vp), h, w_);
let (yo, uo, vo) = icci.forward(&nrm(&yp), &nrm(&up), &nrm(&vp), h, w_);
let den = |p: Vec<f32>| p.into_iter().map(|x| x * 255.0).collect::<Vec<f32>>();
yp = den(yo); up = den(uo); vp = den(vo);
observe(Stage::Icci, &[&yp, &up, &vp]);

60
wai/wai-rs/src/jpegai_icci.rs
View file

@ -16,6 +16,8 @@
// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
#![deny(clippy::disallowed_methods)]
pub use crate::jpegai_synthesis::BadMember;
/// One-level Haar DWT: `[c, h, w]` → `[4c, h/2, w/2]`, channels grouped `LL, HL, LH, HH`.
/// `h` and `w` must be even. Returns the coefficients and the output dims.
pub fn haar_dwt(x: &[f32], c: usize, h: usize, w: usize) -> (Vec<f32>, usize, usize) {
@ -304,40 +306,48 @@ impl IcciNet {
/// Parse the whole eICCI net from the `icci` weight export
/// (`{luma, chroma, conv_last_v, scale_v}`, produced by
/// `tools/wai_jpegai_export_weights.py icci`).
pub fn from_json(v: &serde_json::Value) -> Self {
IcciNet {
luma: IcciPath::from_json(&v["luma"]),
chroma: IcciPath::from_json(&v["chroma"]),
conv_last_v: Conv::from_json(&v["conv_last_v"]),
scale_v: v["scale_v"].as_f64().unwrap() as f32,
}
/// `tools/wai_jpegai_export_weights.py icci`). A member missing or of another type is
/// refused, named by its path ([`BadMember`]); the loader used to panic on one.
pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
Ok(IcciNet {
luma: IcciPath::from_json(&v["luma"]).map_err(|e| e.under("luma"))?,
chroma: IcciPath::from_json(&v["chroma"]).map_err(|e| e.under("chroma"))?,
conv_last_v: Conv::from_json(&v["conv_last_v"]).map_err(|e| e.under("conv_last_v"))?,
scale_v: v["scale_v"].as_f64().ok_or_else(|| BadMember("scale_v".into()))? as f32,
})
}
}
impl Conv {
/// Parse a BN-folded conv from an export node `{w, b, cin, cout, kh, kw, ph, pw}`.
pub fn from_json(v: &serde_json::Value) -> Self {
let g = |k: &str| v[k].as_u64().unwrap() as usize;
let f = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
Conv { w: f("w"), b: f("b"), cin: g("cin"), cout: g("cout"),
kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw") }
/// Parse a BN-folded conv from an export node `{w, b, cin, cout, kh, kw, ph, pw}`. A
/// field missing or of another type is refused ([`BadMember`], named by the field).
pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
let g = |k: &str| v[k].as_u64().and_then(|x| usize::try_from(x).ok()).ok_or_else(|| BadMember(k.into()));
let f = |k: &str| -> Result<Vec<f32>, BadMember> {
v[k].as_array().and_then(|a| a.iter().map(|x| x.as_f64().map(|x| x as f32)).collect::<Option<Vec<f32>>>()).ok_or_else(|| BadMember(k.into()))
};
Ok(Conv { w: f("w")?, b: f("b")?, cin: g("cin")?, cout: g("cout")?,
kh: g("kh")?, kw: g("kw")?, ph: g("ph")?, pw: g("pw")? })
}
}
impl IcciPath {
/// Parse one plane-path (`{conv_first, blocks:[{c1,c2,scale}], conv_last, scale}`).
pub fn from_json(v: &serde_json::Value) -> Self {
IcciPath {
conv_first: Conv::from_json(&v["conv_first"]),
blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock {
conv1: Conv::from_json(&b["c1"]),
conv2: Conv::from_json(&b["c2"]),
scale: b["scale"].as_f64().unwrap() as f32,
}).collect(),
conv_last: Conv::from_json(&v["conv_last"]),
scale: v["scale"].as_f64().unwrap() as f32,
}
pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
let blocks = v["blocks"].as_array().ok_or_else(|| BadMember("blocks".into()))?;
Ok(IcciPath {
conv_first: Conv::from_json(&v["conv_first"]).map_err(|e| e.under("conv_first"))?,
blocks: blocks.iter().enumerate().map(|(i, b)| -> Result<ResBlock, BadMember> {
let at = |e: BadMember| e.under(&format!("blocks.{i}"));
Ok(ResBlock {
conv1: Conv::from_json(&b["c1"]).map_err(|e| at(e.under("c1")))?,
conv2: Conv::from_json(&b["c2"]).map_err(|e| at(e.under("c2")))?,
scale: b["scale"].as_f64().ok_or_else(|| at(BadMember("scale".into())))? as f32,
})
}).collect::<Result<_, _>>()?,
conv_last: Conv::from_json(&v["conv_last"]).map_err(|e| e.under("conv_last"))?,
scale: v["scale"].as_f64().ok_or_else(|| BadMember("scale".into()))? as f32,
})
}
}

145
wai/wai-rs/src/jpegai_ops.rs
View file

@ -1,6 +1,7 @@
//! The integer stages of `wai.image.jpegai` through the integer operator set
//! `wai-ops/1`: the scale path and the two entropy decodes of
//! `extensions/integer-operators.md` §13.6.1, lowered by `jpegai-int/1`.
//! `wai-ops/1`, lowered by `jpegai-int/1`: the scale path and the two entropy decodes of
//! `extensions/integer-operators.md` §13.6.1, and the hyper decoders, context stages and
//! synthesis of §13.6.2 ([`Networks`]).
//!
//! A bundle's scale path is lowered and loaded once ([`load`], [`load_bytes`]): the
//! load checks of §13.6.1 on `norm.json`, the two scale graphs `hs_y` and `hs_uv`
@ -314,6 +315,107 @@ pub fn fac_rows(cdfs_y: &[i64], cdfs_uv: &[i64], p: Plane) -> Vec<Vec<i64>> {
rows
}
// ------------------------------------------------------------ §13.6.2's networks
/// The profile the other integer networks are loaded under (§13.6.2, §11's
/// `role_profiles`): `acc64`. Their layer normalisations, attention and Q16 activations
/// form values above `i32` whatever their weights (§7.3), so no bundle's networks meet
/// `acc32`.
pub(crate) const NETS_PROFILE: Profile = Profile::Acc64;
/// A bundle's other integer networks, lowered and loaded (§13.6.2): `rec_y` and `rec_uv`
/// (the hyper decoder, and for luma the four context stages), and `syn_y` and `syn_uv`
/// (the synthesis).
pub struct Networks {
rec: [Loaded; 2],
syn: [Loaded; 2],
}
impl Networks {
/// The graph of `role` (`rec_y`, `rec_uv`, `syn_y` or `syn_uv`), as loaded.
pub fn graph(&self, role: &str) -> Option<&Loaded> {
match role {
"rec_y" => Some(&self.rec[0]),
"rec_uv" => Some(&self.rec[1]),
"syn_y" => Some(&self.syn[0]),
"syn_uv" => Some(&self.syn[1]),
_ => None,
}
}
/// What the networks hold on the heap: the four graphs.
pub fn heap_bytes(&self) -> u64 {
self.rec.iter().chain(&self.syn).map(|g| g.graph().heap_bytes()).sum()
}
}
/// The networks this crate parsed ([`crate::jpegai_synthesis::Weights`]) in the form
/// `wai-ops` lowers them from: the two hyper decoders, the four context stages and the
/// two synthesis networks.
pub fn nets_of(w: &crate::jpegai_decode::DecoderWeights) -> Result<wai_ops::lower::jpegai_nets::JpegaiNets, Refusal> {
let m = &w.mcm_y;
Ok(wai_ops::lower::jpegai_nets::JpegaiNets {
hyp: [w.luma.hyper_decoder.ptq()?, w.chroma.hyper_decoder.ptq()?],
mcm: [m[0].ptq()?, m[1].ptq()?, m[2].ptq()?, m[3].ptq()?],
synth: [w.luma.synthesis.ptq()?, w.chroma.synthesis.ptq()?],
})
}
/// Lower and load the networks of a parsed bundle under [`NETS_PROFILE`].
pub fn load_networks(w: &crate::jpegai_decode::DecoderWeights) -> Result<Networks, Refusal> {
load_networks_net(nets_of(w)?)
}
/// Lower and load the networks from their files' bytes, in
/// `wai_ops::lower::jpegai_nets::NET_MEMBERS`' order.
pub fn load_networks_bytes(files: [&[u8]; 8]) -> Result<Networks, Refusal> {
load_networks_net(wai_ops::lower::jpegai_nets::parse_nets(files)?)
}
/// The networks of a parsed bundle, loaded under [`NETS_PROFILE`].
pub fn load_networks_net(n: wai_ops::lower::jpegai_nets::JpegaiNets) -> Result<Networks, Refusal> {
let [ry, ruv, sy, suv] = wai_ops::lower::jpegai_nets::nets_loaded(n, NETS_PROFILE)?;
Ok(Networks { rec: [ry, ruv], syn: [sy, suv] })
}
fn run_graph(g: &Loaded, inputs: Vec<Tensor>) -> Result<Tensor, Refusal> {
let mut out = wai_ops::run_with(g, inputs, &[], &mut wai_ops::Sequence::new(), &crate::int_ops::Host)?;
Ok(out.swap_remove(0))
}
/// One plane's reconstructed latent `ŷ [C, 16, 16]` (§13.6.2, `rec_y` or `rec_uv`): the
/// hyper decoder on `z [C, 4, 4]` and, for luma, the four context stages, with the
/// dequantised residual `resi` (Q16, `[C, 16, 16]`). A residual value outside `i32` is
/// refused (`field_range`), and so is any activation outside `i32`.
pub fn reconstruct(n: &Networks, p: Plane, z: Vec<i64>, resi: Vec<i64>) -> Result<Vec<i64>, Refusal> {
let c = p.channels() as u32;
let out = run_graph(&n.rec[p as usize], vec![Tensor::new(c, Z_SIDE, Z_SIDE, z), Tensor::new(c, SIDE, SIDE, resi)])?;
if (out.c, out.h, out.w) != (c, SIDE, SIDE) {
return Err(refuse());
}
Ok(out.data)
}
/// The luma plane `[256, 256]` out of the synthesis, denormalised, in Q16 (§13.6.2,
/// `syn_y`), from the luma latent `[160, 16, 16]` in Q16.
pub fn synthesize_y(n: &Networks, latent_y: Vec<i64>) -> Result<Vec<i64>, Refusal> {
let out = run_graph(&n.syn[0], vec![Tensor::new(160, SIDE, SIDE, latent_y)])?;
if (out.c, out.h, out.w) != (1, 256, 256) {
return Err(refuse());
}
Ok(out.data)
}
/// The chroma planes U then V, each `[256, 256]`, out of the synthesis, denormalised, in
/// Q16 (§13.6.2, `syn_uv`), from the luma latent `[160, 16, 16]` and the chroma latent
/// `[96, 16, 16]` in Q16.
pub fn synthesize_uv(n: &Networks, latent_y: Vec<i64>, latent_uv: Vec<i64>) -> Result<Vec<i64>, Refusal> {
let out = run_graph(&n.syn[1], vec![Tensor::new(160, SIDE, SIDE, latent_y), Tensor::new(96, SIDE, SIDE, latent_uv)])?;
if (out.c, out.h, out.w) != (2, 256, 256) {
return Err(refuse());
}
Ok(out.data)
}
/// The front half as the exact decoder computed it before it decoded through the
/// operator set (test-only, `legacy_int`): the crate's mANS port, its hyper-scale network
/// and the pipeline's derivations. It panics where its port panics, on a stream that runs
@ -415,4 +517,43 @@ mod fast_kernels {
}
}
/// The other integer networks through their lowering (§13.6.2: the committed bundle's
/// four graphs, loaded under `acc64` as `jpegai_ops` loads them): under every fast
/// policy every convolution and transposed convolution of each graph runs on the fast
/// kernel, and every value is the interpreter's, on inputs at small sizes over a
/// range of magnitudes.
#[test]
fn the_jpegai_networks_run_their_convolutions_on_the_fast_kernel() {
let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance/sets/jpegai");
let files: Vec<Vec<u8>> = wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().map(|n| std::fs::read(dir.join(n)).unwrap()).collect();
let nets = load_networks_bytes(std::array::from_fn(|i| files[i].as_slice())).unwrap();
let mut r = Rng(0x0D2C_007B);
for role in ["rec_y", "rec_uv", "syn_y", "syn_uv"] {
let mut ran = 0;
let g = nets.graph(role).unwrap();
let convs = g.graph().ops.iter().filter(|o| matches!(o.p, wai_ops::graph::Params::Conv { .. } | wai_ops::graph::Params::ConvT { .. })).count() as u64;
for k in 0..4 {
let a = 1i64 << (2 + 3 * k);
let ins: Vec<Tensor> = g.graph().inputs.iter().enumerate().map(|(i, rec)| {
let c = g.graph().tensors[rec.tensor as usize].channels;
let (h, w) = if role.starts_with("rec") { if i == 0 { (1, 1) } else { (4, 4) } } else { (2, 2) };
let n = (c * h * w) as usize;
let d = if role.starts_with("rec") && i == 0 { (0..n).map(|_| r.below(64) as i64 - 31).collect() } else { r.vec(n, a, 10) };
Tensor::new(c, h, w, d)
}).collect();
let run = || wai_ops::run_with(g, ins.clone(), &[], &mut wai_ops::Sequence::new(), &crate::int_ops::Host);
let (want, none) = fast_conv_runs(|| with_policy(Policy::reference(), run));
assert_eq!(none, 0, "{role}: the reference policy runs no fast kernel");
for p in FAST {
let (got, runs) = fast_conv_runs(|| with_policy(p, run));
assert!(got == want, "{role} {p:?}: the outputs differ from the interpreter alone");
if want.is_ok() {
assert_eq!(runs, convs, "{role} {p:?}: every convolution runs on the fast kernel");
ran += 1;
}
}
}
assert!(ran > 0, "{role}: some decode completes, so its every operator ran");
}
}
}

9
wai/wai-rs/src/jpegai_pipeline.rs
View file

@ -27,7 +27,8 @@
//! between YUV and RGB (BT.709 full-range, data normalised to [0,1]) — the last pixel
//! transform before the output image.
//!
//! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine:
//! **Latent reconstruction.** [`reconstruct_yhat`] (test-only, `legacy_int`: the decode
//! runs it as the graph `rec_y` of integer-operators §13.6.2) is the SGMM autoregressive spine:
//! it turns the dequantized residual back into the synthesis-input latent `y_hat` by
//! running the four MCM context stages over a 2×2 checkerboard, each stage predicting
//! a mean from the hyper params plus the already-reconstructed stages.
@ -36,15 +37,18 @@
#![deny(clippy::disallowed_methods)]
use crate::float_repro::{exp, round_half_away, round_half_even};
#[cfg(any(test, feature = "legacy_int"))]
use crate::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights};
/// The 2×2 sub-lattice each MCM stage owns, in reference stage order:
/// stage 0 → (0,0), stage 1 → (1,1), stage 2 → (0,1), stage 3 → (1,0).
/// (The reference's `down_shuffle` returns `(k00, k11, k01, k10)`; `up_shuffle` is
/// its exact inverse. The diagonal pair is decoded first — a checkerboard.)
#[cfg(any(test, feature = "legacy_int"))]
const SUBPOS: [(usize, usize); 4] = [(0, 0), (1, 1), (0, 1), (1, 0)];
/// Split `[c, 2h, 2w]` into the four `[c, h, w]` sub-lattices (reference `down_shuffle`).
#[cfg(any(test, feature = "legacy_int"))]
fn down_shuffle(y: &[i64], c: usize, h2: usize, w2: usize) -> [Vec<i64>; 4] {
let (h, w) = (h2 / 2, w2 / 2);
let mut parts = std::array::from_fn(|_| vec![0i64; c * h * w]);
@ -61,6 +65,7 @@ fn down_shuffle(y: &[i64], c: usize, h2: usize, w2: usize) -> [Vec<i64>; 4] {
}
/// Interleave the four `[c, h, w]` sub-lattices back into `[c, 2h, 2w]` (`up_shuffle`).
#[cfg(any(test, feature = "legacy_int"))]
fn up_shuffle(parts: &[Vec<i64>; 4], c: usize, h: usize, w: usize) -> Vec<i64> {
let (h2, w2) = (h * 2, w * 2);
let mut out = vec![0i64; c * h2 * w2];
@ -95,6 +100,7 @@ fn up_shuffle(parts: &[Vec<i64>; 4], c: usize, h: usize, w: usize) -> Vec<i64> {
/// stage-0 model ignores them). The context concatenation is **ascending** (`y_0`
/// first) — the reference does `cat(y_hat_list)`, despite the ONNX input being
/// labelled `y[1-0]`.
#[cfg(any(test, feature = "legacy_int"))]
pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize, mcm: [&Weights; 4]) -> Vec<i64> {
let (h, w) = (h2 / 2, w2 / 2);
let plane = c * h * w;
@ -134,6 +140,7 @@ pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usiz
/// is deliberate: the two paths are structurally different, not parameterised.
///
/// `resi` is `[c, 2h, 2w]` (Q16); `psi` is `[4c, h, w]` (Q16). Returns `[c, 2h, 2w]`.
#[cfg(any(test, feature = "legacy_int"))]
pub fn reconstruct_yhat_uv(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize) -> Vec<i64> {
let (h, w) = (h2 / 2, w2 / 2);
let plane = c * h * w;

139
wai/wai-rs/src/jpegai_synthesis.rs
View file

@ -1,3 +1,8 @@
//! The weight stores of `wai.image.jpegai`'s integer networks ([`Weights`], [`IntNet`])
//! and, test-only (`legacy_int`), the direct networks the exact decoder ran before they
//! ran through the integer operator set (`crate::jpegai_ops`, integer-operators §13.6):
//! the reference the golden-equality tests and the differential hold the graphs to.
//!
//! Whole-net **byte-exact integer JPEG-AI synthesis** (model_y + model_uv, tools_0)
//! — threads the real Restormer topology (traced from the ISO reference ONNX)
//! through the integer ops in [`crate::int_transformer`] + [`crate::int_hyper_synth`],
@ -30,13 +35,43 @@
// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
#![deny(clippy::disallowed_methods)]
#[cfg(any(test, feature = "legacy_int"))]
use crate::int_kernels::{par_map, par_zip};
#[cfg(any(test, feature = "legacy_int"))]
use crate::int_transformer as it;
use std::collections::HashMap;
#[cfg(any(test, feature = "legacy_int"))]
const F: u32 = 16;
#[cfg(any(test, feature = "legacy_int"))]
const S: i64 = 1 << F;
/// A member of a weight file that is not of its kind's form, by its name (`""` for the
/// file itself, when it is not a JSON object). The bundle is refused with
/// `model_mismatch` (`extensions/integer-operators.md` §13.6.2, *The members*); the
/// loaders used to panic on one.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BadMember(pub String);
impl std::fmt::Display for BadMember {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if self.0.is_empty() {
write!(f, "model_mismatch: a weight file is not a JSON object")
} else {
write!(f, "model_mismatch: the weight member `{}` is not of its kind's form", self.0)
}
}
}
impl std::error::Error for BadMember {}
impl BadMember {
/// The same member, named under `parent` (`parent.name`).
pub fn under(self, parent: &str) -> BadMember {
BadMember(if self.0.is_empty() { parent.to_string() } else { format!("{parent}.{}", self.0) })
}
}
/// A sink-supplied weight store: the PTQ export JSON (`{name: {kind, dims, ...}}`)
/// parsed into typed tensors. `conv`/`q16` fetch by full initializer name.
pub struct Weights {
@ -45,35 +80,47 @@ pub struct Weights {
}
impl Weights {
/// Parse the exported weights JSON.
pub fn from_json(v: &serde_json::Value) -> Self {
/// Parse the exported weights JSON: each member of kind `conv` with integer `dims`,
/// integer weights `w`, an integer `mult` and a non-negative `shift`, and each of kind
/// `q16` with integer values `v` (§13.6.2, *The members*); members of another kind are
/// ignored. A member of either kind with a field missing or of another type is refused
/// ([`BadMember`]), and so is a file that is not an object. (The ranges and the shapes
/// are the lowering's to check.)
pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
let mut conv = HashMap::new();
let mut q16 = HashMap::new();
for (k, e) in v.as_object().expect("weights root object") {
match e["kind"].as_str() {
let ints = |x: &serde_json::Value| -> Option<Vec<i64>> { x.as_array()?.iter().map(serde_json::Value::as_i64).collect() };
for (k, e) in v.as_object().ok_or_else(|| BadMember(String::new()))? {
let bad = || BadMember(k.clone());
match e.get("kind").and_then(serde_json::Value::as_str) {
Some("conv") => {
let w = e["w"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
let dims = e["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
conv.insert(k.clone(), (w, dims, e["mult"].as_i64().unwrap(), e["shift"].as_u64().unwrap() as u32));
let w = ints(&e["w"]).ok_or_else(bad)?;
let dims = e["dims"].as_array().and_then(|a| a.iter().map(|x| x.as_u64().and_then(|d| usize::try_from(d).ok())).collect::<Option<Vec<usize>>>()).ok_or_else(bad)?;
let mult = e["mult"].as_i64().ok_or_else(bad)?;
let shift = e["shift"].as_u64().and_then(|s| u32::try_from(s).ok()).ok_or_else(bad)?;
conv.insert(k.clone(), (w, dims, mult, shift));
}
Some("q16") => {
q16.insert(k.clone(), e["v"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect());
q16.insert(k.clone(), ints(&e["v"]).ok_or_else(bad)?);
}
_ => {}
}
}
Weights { conv, q16 }
Ok(Weights { conv, q16 })
}
/// Output-channel count of a conv, read from its weight dims. Lets the nets be
/// driven by their own widths instead of a hardcoded plane size.
#[cfg(any(test, feature = "legacy_int"))]
pub fn out_channels(&self, module: &str) -> usize {
self.conv.get(&format!("{module}.weight"))
.unwrap_or_else(|| panic!("missing conv {module}")).1[0]
}
#[cfg(any(test, feature = "legacy_int"))]
fn bias(&self, module: &str, cout: usize) -> Vec<i64> {
self.q16.get(&format!("{module}.bias")).cloned().unwrap_or_else(|| vec![0; cout])
}
/// grouped conv by module name (weight `<module>.weight`, bias `<module>.bias`).
#[cfg(any(test, feature = "legacy_int"))]
fn conv(&self, m: &str, x: &[i64], cin: usize, h: usize, w: usize, k: usize, stride: usize, pad: usize, groups: usize) -> (Vec<i64>, usize, usize) {
let (wt, dims, mult, shift) = self.conv.get(&format!("{m}.weight")).unwrap_or_else(|| panic!("missing conv {m}"));
let cout = dims[0];
@ -82,6 +129,7 @@ impl Weights {
let (o, ho, wo) = it::conv_grouped(x, cin, h, w, wt, cout, k, stride, pad, groups, &m_v, &b, *shift);
(o, ho, wo)
}
#[cfg(any(test, feature = "legacy_int"))]
fn convt(&self, m: &str, x: &[i64], cin: usize, h: usize, w: usize, k: usize, stride: usize, pad: usize, opad: usize) -> (Vec<i64>, usize, usize) {
let (wt, dims, mult, shift) = self.conv.get(&format!("{m}.weight")).unwrap_or_else(|| panic!("missing convT {m}"));
let cout = dims[1]; // ConvTranspose weight [cin, cout, k, k]
@ -89,21 +137,47 @@ impl Weights {
let b = self.bias(m, cout);
it::conv_transpose(x, cin, h, w, wt, cout, k, stride, pad, opad, &m_v, &b, *shift)
}
#[cfg(any(test, feature = "legacy_int"))]
fn q16v(&self, name: &str) -> Vec<i64> {
self.q16.get(name).cloned().unwrap_or_else(|| panic!("missing q16 {name}"))
}
/// The file's members of kind `conv` and `q16`, in the form `wai-ops` lowers them
/// from (`extensions/integer-operators.md` §13.6.2, *The members*): the lowering of
/// the file's bytes reads these same members. A shift past 127, or weights of another
/// count than their dims give, which no such member holds, is refused
/// (`model_mismatch`). Two members of one name, which the bytes' lowering refuses,
/// this parsed form cannot hold.
pub fn ptq(&self) -> Result<wai_ops::lower::jpegai_nets::PtqNet, wai_ops::Refusal> {
use wai_ops::lower::jpegai_nets::{PtqMember, PtqNet};
let refuse = || wai_ops::Refusal::new(wai_ops::Code::ModelMismatch, wai_ops::Site::Lowering);
let mut net = PtqNet::default();
for (k, (w, dims, mult, shift)) in &self.conv {
let shift = u8::try_from(*shift).ok().filter(|&s| s <= 127).ok_or_else(refuse)?;
if dims.iter().try_fold(1usize, |a, &d| a.checked_mul(d)) != Some(w.len()) {
return Err(refuse());
}
net.insert(k.clone(), PtqMember::Conv { dims: dims.iter().map(|&d| d as u64).collect(), w: w.clone(), mult: *mult, shift })?;
}
for (k, v) in &self.q16 {
net.insert(k.clone(), PtqMember::Q16 { v: v.clone() })?;
}
Ok(net)
}
}
// The element-wise helpers compute each output element with the same pure
// function of its own inputs whatever the partition; `par_map`/`par_zip` only spread
// the elements over workers (one worker under the reference policy).
#[cfg(any(test, feature = "legacy_int"))]
fn relu(x: &[i64]) -> Vec<i64> {
par_map(x, |v| if v < 0 { 0 } else { v })
}
#[cfg(any(test, feature = "legacy_int"))]
fn relu6_of(x: &[i64]) -> Vec<i64> {
par_map(x, |v| it::clip(v, 0, 6 * S))
}
/// conv1 → ReLU → conv2 → + input residual block (`first_stage.0.0`, `CAB.*`).
#[cfg(any(test, feature = "legacy_int"))]
fn res_block(p: &str, x: &[i64], c: usize, h: usize, w: usize, wt: &Weights) -> Vec<i64> {
let (a, _, _) = wt.conv(&format!("{p}.conv1"), x, c, h, w, 3, 1, 1, 1);
let ar = relu(&a);
@ -112,6 +186,7 @@ fn res_block(p: &str, x: &[i64], c: usize, h: usize, w: usize, wt: &Weights) ->
}
/// DepthToSpace **CRD** mode (PyTorch PixelShuffle): the net uses CRD, not the DCR
/// in `int_transformer::depth_to_space`. `out[cc, h·b+bh, w·b+bw] = in[cc·b² + bh·b + bw, h, w]`.
#[cfg(any(test, feature = "legacy_int"))]
fn depth_to_space_crd(x: &[i64], c: usize, h: usize, w: usize, b: usize) -> (Vec<i64>, usize, usize, usize) {
let co = c / (b * b);
let (ho, wo) = (h * b, w * b);
@ -130,20 +205,24 @@ fn depth_to_space_crd(x: &[i64], c: usize, h: usize, w: usize, b: usize) -> (Vec
}
(out, co, ho, wo)
}
#[cfg(any(test, feature = "legacy_int"))]
fn emul(a: &[i64], b: &[i64]) -> Vec<i64> {
par_zip(a, b, |x, y| ((x as i128 * y as i128) >> F) as i64)
}
#[cfg(any(test, feature = "legacy_int"))]
fn eadd(a: &[i64], b: &[i64]) -> Vec<i64> {
par_zip(a, b, |x, y| x + y)
}
/// Residual multiplicative gate `x · (f(x) + 1)` (the exact form: the ONNX gate
/// `Add`s carry const 1.0 with no activation). Near-identity when f≈0.
#[cfg(any(test, feature = "legacy_int"))]
fn gate(x: &[i64], modulation: &[i64]) -> Vec<i64> {
let one_plus: Vec<i64> = par_map(modulation, |v| v + S);
emul(x, &one_plus)
}
/// One TAB (Restormer block) via the merged executors, weights by module prefix.
#[cfg(any(test, feature = "legacy_int"))]
fn tab(p: &str, x: &[i64], c: usize, heads: usize, h: usize, w: usize, wt: &Weights, eps: u128) -> Vec<i64> {
let attn = it::MdtaBlockWeights {
c, heads, h, w,
@ -187,29 +266,35 @@ fn tab(p: &str, x: &[i64], c: usize, heads: usize, h: usize, w: usize, wt: &Weig
/// in Q16. The reference post-filters operate on this **float** value (0..255) and round
/// to u8 only after BT.709, so [`decode_y`]/[`decode_uv`] round here for a standalone
/// plane while [`decode_y_f32`]/[`decode_uv_f32`] keep the float for the filter chain.
#[cfg(any(test, feature = "legacy_int"))]
fn denorm_q16(v: i64) -> i64 {
let half = S / 2;
(v * 128 + 127 * S + half).clamp(0, 255 * S)
}
#[cfg(any(test, feature = "legacy_int"))]
fn denorm_u8(planes: &[i64]) -> Vec<u8> {
let half = S / 2;
planes.iter().map(|&v| ((denorm_q16(v) + half) >> F).clamp(0, 255) as u8).collect()
}
#[cfg(any(test, feature = "legacy_int"))]
fn denorm_f32(planes: &[i64]) -> Vec<f32> {
planes.iter().map(|&v| denorm_q16(v) as f32 / S as f32).collect()
}
#[cfg(any(test, feature = "legacy_int"))]
pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec<u8> {
denorm_u8(&decode_y_planes(latent, wt))
}
/// The luma plane as **float** samples in `[0,255]` (pre-round) — the domain the
/// post-filter chain runs in.
#[cfg(any(test, feature = "legacy_int"))]
pub fn decode_y_f32(latent: &[i64], wt: &Weights) -> Vec<f32> {
denorm_f32(&decode_y_planes(latent, wt))
}
/// The luma synthesis up to (but not including) the output denormalize: the raw
/// `conv4_t` output `[1,256,256]` in Q16.
#[cfg(any(test, feature = "legacy_int"))]
fn decode_y_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5
// first_stage.0.0 residual block: conv1→ReLU→conv2 + latent [160,16,16]
@ -278,17 +363,20 @@ fn decode_y_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
/// • TAM: the two TABs run **at the full DepthToSpace resolution** — no ds/us conv;
/// • tail: `conv4_t(64→8, ⇑2)` → `DepthToSpace(CRD,2)` → denorm → **2-channel UV**.
/// Widths are 64 (vs model_y's 128). CAB / iact gating / TAB / denorm are identical.
#[cfg(any(test, feature = "legacy_int"))]
pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec<u8> {
denorm_u8(&decode_uv_planes(latent, wt))
}
/// The two chroma planes as **float** samples in `[0,255]` (pre-round) — the domain the
/// post-filter chain runs in.
#[cfg(any(test, feature = "legacy_int"))]
pub fn decode_uv_f32(latent: &[i64], wt: &Weights) -> Vec<f32> {
denorm_f32(&decode_uv_planes(latent, wt))
}
/// The chroma synthesis up to (but not including) the output denormalize: the raw
/// depth-to-space output `[2,256,256]` in Q16.
#[cfg(any(test, feature = "legacy_int"))]
fn decode_uv_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5
let (h0, w0) = (16usize, 16usize);
@ -357,10 +445,12 @@ fn decode_uv_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
/// Lift entropy-decoded integer symbols into the Q16 fixed-point the PTQ nets expect.
/// Named rather than inline so the scaling is visible at the call site — see the
/// failure mode described on [`decode_hyper`].
#[cfg(any(test, feature = "legacy_int"))]
pub fn q16_from_symbols(symbols: &[i64]) -> Vec<i64> {
symbols.iter().map(|&v| v * S).collect()
}
#[cfg(any(test, feature = "legacy_int"))]
pub fn decode_hyper(z: &[i64], cin: usize, h: usize, w: usize, wt: &Weights) -> (Vec<i64>, usize, usize) {
// Widths come from the weights: chroma is the same topology at 96 / 384 where luma
// is 160 / 640, so nothing here may assume the luma channel count.
@ -381,14 +471,21 @@ pub struct IntNet {
t: HashMap<String, (Vec<i64>, Vec<usize>)>,
}
impl IntNet {
pub fn from_json(v: &serde_json::Value) -> Self {
/// Parse a scale network's file: every member that is a tensor `{ "dims", "v" }` of
/// integers. A member of another form is not kept, so a member §13.6.1 reads that is
/// not of its form is missing, and the scale path's load refuses it, as the lowering
/// of the file's bytes does (`model_mismatch`). A file that is not an object is
/// refused ([`BadMember`]).
pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
let mut t = HashMap::new();
for (k, e) in v.as_object().expect("intnet root object") {
let vals = e["v"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
let dims = e["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
t.insert(k.clone(), (vals, dims));
for (k, e) in v.as_object().ok_or_else(|| BadMember(String::new()))? {
let vals = e.get("v").and_then(serde_json::Value::as_array).and_then(|a| a.iter().map(serde_json::Value::as_i64).collect::<Option<Vec<i64>>>());
let dims = e.get("dims").and_then(serde_json::Value::as_array).and_then(|a| a.iter().map(|x| x.as_u64().and_then(|d| usize::try_from(d).ok())).collect::<Option<Vec<usize>>>());
if let (Some(vals), Some(dims)) = (vals, dims) {
t.insert(k.clone(), (vals, dims));
}
}
IntNet { t }
Ok(IntNet { t })
}
#[cfg(any(test, feature = "legacy_int"))]
fn get(&self, n: &str) -> &(Vec<i64>, Vec<usize>) {
@ -482,6 +579,7 @@ pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNe
/// **MCM stage 0** (common_modules/MCM) — the multi-stage context model's first
/// stage: `p0 → conv1 → ReLU → conv2 → ReLU → conv3 → y[0]` (bias-free 1×1 convs,
/// width 160). Q16 fixed-point; `wt` = this stage's own PTQ weights.
#[cfg(any(test, feature = "legacy_int"))]
pub fn mcm_stage0(p0: &[i64], h: usize, w: usize, wt: &Weights) -> Vec<i64> {
let (x, _, _) = wt.conv("fusion_pred_net.conv1", p0, 160, h, w, 1, 1, 0, 1);
let x = relu(&x);
@ -494,6 +592,7 @@ pub fn mcm_stage0(p0: &[i64], h: usize, w: usize, wt: &Weights) -> Vec<i64> {
/// `cy = conv.1(conv.0(ctx))` (conv.0 1×1 `ctx_ch→160`, conv.1 3×3 grouped-5 +bias),
/// then `Concat(cy, p)` [→320] → `fusion.conv1 → ReLU → conv2 → ReLU → conv3 → y[i]`.
/// `ctx` is the caller-assembled context `concat(y[i-1..0])` (160/320/480 ch).
#[cfg(any(test, feature = "legacy_int"))]
pub fn mcm_stage_n(p: &[i64], ctx: &[i64], ctx_ch: usize, h: usize, w: usize, wt: &Weights) -> Vec<i64> {
let (c0, _, _) = wt.conv("conv.0", ctx, ctx_ch, h, w, 1, 1, 0, 1); // ctx_ch→160, 1×1
let (cy, _, _) = wt.conv("conv.1", &c0, 160, h, w, 3, 1, 1, 5); // 3×3 grouped-5, +bias
@ -522,7 +621,7 @@ mod tests {
};
if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
// deterministic synthetic latent [160,16,16] in Q16 (a ramp, non-degenerate)
let latent: Vec<i64> = (0..160 * 16 * 16).map(|i| ((i as i64 % 97) - 48) * (S / 64)).collect();
let y = decode_y(&latent, &wt);
@ -543,7 +642,7 @@ mod tests {
};
if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
// deterministic synthetic 256-ch latent [256,16,16] in Q16 (a ramp)
let latent: Vec<i64> = (0..256 * 16 * 16).map(|i| ((i as i64 % 97) - 48) * (S / 64)).collect();
let uv = decode_uv(&latent, &wt);
@ -565,7 +664,7 @@ mod tests {
};
if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
let (c, h, w) = (160usize, 8usize, 8usize);
// integer-valued hyper-latent (entropy-decoded symbols) in Q16
let z: Vec<i64> = (0..c * h * w).map(|i| ((i as i64 % 11) - 5) * S).collect();
@ -588,7 +687,7 @@ mod tests {
};
if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
let net = IntNet::from_json(&v);
let net = IntNet::from_json(&v).unwrap();
let (c, h, w) = (160usize, 8usize, 8usize);
let z: Vec<i64> = (0..c * h * w).map(|i| ((i as i64 * 37) % 291) - 140).collect();
let (is, co, ho, wo) = decode_hyper_scale(&z, c, h, w, &net);
@ -612,7 +711,7 @@ mod tests {
let p = format!("{dir}/mcm_stage{i}_y.json");
if !std::path::Path::new(&p).exists() { return None; }
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&p).unwrap()).unwrap();
Some(Weights::from_json(&v))
Some(Weights::from_json(&v).unwrap())
};
let (Some(w0), Some(w1)) = (load(0), load(1)) else { eprintln!("skip: MCM weights not present"); return; };
let (h, w) = (16usize, 16usize);

9
wai/wai-rs/src/lib.rs
View file

@ -112,7 +112,9 @@ pub mod int_mlicpp;
// L2-norm / ELU / Sigmoid / MDTA) — the pure-i64 arithmetic a learned ATTENTION
// image codec (JPEG-AI's Restormer synthesis) needs beyond the conv/IGDN stack
// above. NOTE: distinct from `int_transform` (a signal transform). Pure i64.
#[cfg(feature = "neural_int")]
// Test-only (`legacy_int`): the direct networks of `jpegai_synthesis`, which the decode
// runs through the integer operator set (integer-operators §13.6.2), are built on it.
#[cfg(all(feature = "neural_int", any(test, feature = "legacy_int")))]
pub mod int_transformer;
// Whole-net byte-exact integer JPEG-AI synthesis (Restormer, model_y) — threads the
@ -152,8 +154,9 @@ pub mod float_repro;
#[cfg(feature = "neural_int")]
pub mod jpegai_decode;
// JPEG-AI's scale path and its two entropy decodes through the integer operator set
// (`extensions/integer-operators.md` §13.6.1, lowered by `jpegai-int/1`).
// JPEG-AI's integer stages through the integer operator set
// (`extensions/integer-operators.md` §13.6, lowered by `jpegai-int/1`): the scale path and
// its two entropy decodes, and the hyper decoders, context stages and synthesis.
#[cfg(feature = "neural_int")]
pub mod jpegai_ops;

77
wai/wai-rs/src/neural/mod.rs
View file

@ -1014,15 +1014,22 @@ enum Exact {
/// process's store of loaded sets (`int_set_cache`), keyed by the SHA-256 of their
/// listing; a bundle that fails a load check of integer-operators §13.6.1, or whose
/// scale graphs do not meet `acc32`, is refused (`model_mismatch`).
/// - The hyper decoders, context stages and synthesis (`weights/hyp_{y,uv}.json`,
/// `weights/mcm{0..3}_y.json`, `weights/synth_{y,uv}.json`) are lowered and loaded from
/// those members' bytes (`jpegai_ops::load_networks_bytes`, integer-operators §13.6.2,
/// under `acc64`), through the same store, keyed by their listing; a weight file that
/// fails §13.6.2's load checks, or networks that do not meet `acc64`, are refused
/// (`model_mismatch`).
/// - A member of the bundle that the parameters lack is `set_incomplete`; a member
/// that is not JSON, or a `norm.json` without the tables the float stages read, is
/// `model_mismatch`. (The float stages' weight loaders still panic on a member whose
/// fields are mistyped; those networks are not yet lowered.)
/// - A substream the entropy decode cannot read is refused (`stream_exhausted`).
/// that is not JSON, an eICCI net whose fields are not of their types, or a
/// `norm.json` without the tables the float stages read, is `model_mismatch`. No
/// loader panics on a mistyped member.
/// - A substream the entropy decode cannot read is refused (`stream_exhausted`), and a
/// dequantised residual, a latent or an activation of the networks past i32
/// (`field_range`).
fn jpegai_exact(payload: &[u8], params: &ParamSource) -> Exact {
use crate::jpegai_decode::{decode_loaded, DecoderWeights, NormTables, PlaneWeights, Refusal};
use crate::jpegai_decode::{decode_loaded, NormTables, Refusal};
use crate::jpegai_icci::IcciNet;
use crate::jpegai_synthesis::{IntNet, Weights};
use crate::int_set_cache::{LoadedSets, SetKey};
let Some(norm_b) = params.try_sibling("norm.json") else { return Exact::Unavailable };
@ -1051,25 +1058,28 @@ fn jpegai_exact(payload: &[u8], params: &ParamSource) -> Exact {
Err(r) => return refused(r.code.as_str(), "the bundle's scale path fails a load check of integer-operators §13.6.1, or does not meet acc32".into()),
};
let json = |name: &str| -> Option<serde_json::Value> { serde_json::from_slice(&bytes[name]).ok() };
let mut parsed = std::collections::BTreeMap::new();
for m in MEMBERS {
let name = format!("weights/{m}.json");
match json(&name) {
Some(v) => { parsed.insert(m, v); }
None => return refused("model_mismatch", format!("{name} is not JSON")),
}
}
// The other integer networks (§13.6.2), from their eight files' bytes.
let nets_files: [&[u8]; 8] = wai_ops::lower::jpegai_nets::NET_MEMBERS.map(|n| &*bytes[n]);
let nets_members: std::collections::BTreeMap<String, &[u8]> =
wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().zip(nets_files).map(|(n, b)| (n.to_string(), b)).collect();
let nets = LoadedSets::global().get_or_load(
|| SetKey::of(crate::codecs::CAP_IMAGE_JPEGAI, crate::jpegai_ops::LOWERING, crate::jpegai_ops::NETS_PROFILE, &nets_members),
|| crate::jpegai_ops::load_networks_bytes(nets_files),
);
let nets = match nets {
Ok(n) => n,
Err(r) => return refused(r.code.as_str(), "the bundle's networks fail a load check of integer-operators §13.6.2, or do not meet acc64".into()),
};
// The eICCI net, a float stage, parsed as its loader reads it.
let icci = match serde_json::from_slice::<serde_json::Value>(&bytes["weights/icci.json"]).ok().map(|v| IcciNet::from_json(&v)) {
Some(Ok(n)) => n,
Some(Err(e)) => return refused("model_mismatch", format!("weights/icci.json: {e}")),
None => return refused("model_mismatch", "weights/icci.json is not JSON".into()),
};
let Ok(nv) = serde_json::from_slice::<serde_json::Value>(&norm_b) else {
return refused("model_mismatch", "norm.json is not JSON".into());
};
let w = |m: &str| Weights::from_json(&parsed[m]);
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: IntNet::from_json(&parsed["hs_y"]), hyper_decoder: w("hyp_y"), synthesis: w("synth_y") },
chroma: PlaneWeights { hyper_scale: IntNet::from_json(&parsed["hs_uv"]), hyper_decoder: w("hyp_uv"), synthesis: w("synth_uv") },
mcm_y: [w("mcm0_y"), w("mcm1_y"), w("mcm2_y"), w("mcm3_y")],
icci: IcciNet::from_json(&parsed["icci"]),
};
// The tables the float stages read, beside the ones the scale path loaded.
let ints = |k: &str| -> Option<Vec<i64>> { nv.get(k)?.as_array()?.iter().map(serde_json::Value::as_i64).collect() };
@ -1095,7 +1105,7 @@ fn jpegai_exact(payload: &[u8], params: &ParamSource) -> Exact {
return refused("model_mismatch", "norm.json lacks a table the float stages read".into());
};
match decode_loaded(payload, &weights, &norm, &sp, &mut |_, _| {}) {
match decode_loaded(payload, &icci, &norm, &sp, &nets, &mut |_, _| {}) {
Ok(rgb) => Exact::Decoded(image::DecodedImage { rgb: rgb.data, width: rgb.w as u32, height: rgb.h as u32 }),
Err(Refusal::Scope(_)) => Exact::Unavailable,
Err(r) => Exact::Refused(DecodeError::InvalidPayload(format!("wai.image.jpegai: {r}"))),
@ -1132,8 +1142,9 @@ mod jpegai_exact_dispatch_tests {
assert!(matches!(super::jpegai_exact(&[], &super::ParamSource::Unpinned("/nonexistent/wai/model/dir".into())), super::Exact::Unavailable));
}
/// A refused exact decode is the decode's result (integer-operators §13.6.1): the
/// bundle, with one bound of 256 in `norm.json`, is refused `model_mismatch`, and the
/// A refused exact decode is the decode's result (integer-operators §13.6): the
/// bundle, with one bound of 256 in `norm.json`, a network weight that is not an
/// integer, or an eICCI field of another type, is refused `model_mismatch`, and the
/// dispatch reports that refusal rather than running the float decode; a bundle
/// missing a member is `set_incomplete`. The intact bundle decodes. Needs the bundle
/// and the codestream (`WAI_JPEGAI_MODEL_DIR`, `WAI_JPEGAI_CODESTREAM`).
@ -1163,5 +1174,21 @@ mod jpegai_exact_dispatch_tests {
let mut short = files.clone();
short.remove("weights/mcm2_y.json");
assert!(matches!(decode(&short), Err(super::DecodeError::PriorMismatch(super::PinRefusal::SetIncomplete { .. }))), "a member missing");
// A network's weight that is not an integer (integer-operators §13.6.2), and an
// eICCI field of another type: each refused model_mismatch, where the loaders
// panicked, and not decoded through the float path.
for (file, edit) in [
("weights/synth_y.json", (|v: &mut serde_json::Value| v["conv2_t.weight"]["w"][0] = serde_json::json!(0.5)) as fn(&mut serde_json::Value)),
("weights/icci.json", |v: &mut serde_json::Value| v["scale_v"] = serde_json::json!("1")),
] {
let mut bad = files.clone();
let mut v: serde_json::Value = serde_json::from_slice(&files[file]).unwrap();
edit(&mut v);
bad.insert(file.into(), serde_json::to_vec(&v).unwrap());
match decode(&bad) {
Err(super::DecodeError::InvalidPayload(m)) => assert!(m.starts_with("wai.image.jpegai: model_mismatch"), "{file}: {m}"),
other => panic!("{file}: a refused exact decode must be the result: {:?}", other.map(|_| ())),
}
}
}
}

2
wai/wai-rs/tests/int_kernels_geometry.rs
View file

@ -24,7 +24,7 @@
//! through that graph under every policy (`int_audio` unit tests). CI runs this file
//! beside `int_kernels_vectors` on every `wai-conformance` leg, with overflow checks
//! on and off.
#![cfg(feature = "neural_int")]
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
use serde_json::json;
use wai::int_kernels::{with_policy, Policy};

324
wai/wai-rs/tests/int_ops_differential.rs
View file

@ -12,7 +12,9 @@
//! `wai.neural.int_mlicpp` (`wmb1/1`, its three networks); and `wai.world.model`'s integer
//! renderer (`world/1`, informative: `T` and `g_s`), against its direct renderer with the
//! direct guard every other decoder ran; and `wai.image.jpegai`'s scale path and entropy
//! decodes (`jpegai-int/1`), against its direct computation; and `WIV1` and `WIR1`
//! decodes (`jpegai-int/1`), against its direct computation, and its hyper decoders,
//! context stages and synthesis (the graphs of §13.6.2), against its direct networks; and
//! `WIV1` and `WIR1`
//! (`wai.video.int_motion` and the P-frames of `wai.video.int_hyper`, and the four
//! `int_refine` capabilities), through the stage operators alone, which make no planned
//! break (with `video_int` and `staged`).
@ -2436,7 +2438,7 @@ fn jpegai_planned_break(c: &JCase) -> Option<&'static str> {
}
}
let b = |v: &Value| serde_json::to_vec(v).unwrap();
let g = wai_ops::lower::jpegai(&b(n), &b(&c.hs[0]), &b(&c.hs[1])).expect("a bundle that passes the load checks lowers");
let g = wai_ops::lower::jpegai_scale(&b(n), &b(&c.hs[0]), &b(&c.hs[1])).expect("a bundle that passes the load checks lowers");
if g.roles.iter().any(|(_, g)| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc32).is_err()) {
assert!(g.roles.iter().all(|(_, g)| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64).is_ok()), "the generator's sums fit i64");
return Some("acc32");
@ -2460,7 +2462,7 @@ fn jpegai_parsed(c: &JCase) -> (wai::jpegai_decode::NormTables, [wai::jpegai_syn
lef_thr: [0; 3],
lef_mag: [0.0; 3],
};
(n, [wai::jpegai_synthesis::IntNet::from_json(&c.hs[0]), wai::jpegai_synthesis::IntNet::from_json(&c.hs[1])])
(n, [wai::jpegai_synthesis::IntNet::from_json(&c.hs[0]).unwrap(), wai::jpegai_synthesis::IntNet::from_json(&c.hs[1]).unwrap()])
}
/// The committed test codestream's SOZ, SORP and SORS payloads.
@ -2867,3 +2869,319 @@ fn wir1_random_layers_apply_alike() {
assert!(alike * 2 > n, "the generator must produce decodable units: {alike} of {n}");
assert!(outcomes.len() >= 4, "too few kinds of outcome: {outcomes:?}");
}
// ------------------------------------------------- jpegai's networks (M7b, §13.6.2)
//
// `wai.image.jpegai`'s hyper decoders, context stages and synthesis (integer-operators
// §13.6.2, lowered by `jpegai-int/1` to `rec_y`, `rec_uv`, `syn_y` and `syn_uv`) against
// the direct networks (`jpegai_synthesis`, `jpegai_pipeline`), on random networks with
// the shapes the direct networks read and random inputs. The direct networks refuse
// nothing. Where the operator set refuses, the refusal is a planned break: a network file
// that fails a load check (`model_mismatch`, from the files alone), or a value past `A32`
// (`field_range`) at a checked tensor, which the direct networks computed on.
/// A random weight file in the export's form, member by member.
struct NetFile(serde_json::Map<String, Value>);
impl NetFile {
fn new() -> NetFile {
NetFile(serde_json::Map::new())
}
/// `m.weight` of `dims` with `taps` products per output, its multiplier scaled so that
/// the convolution keeps its input's magnitude, and `m.bias` of `bias` values when
/// given. `hot` makes one convolution in about 64 grow its input past `A32`.
fn conv(&mut self, r: &mut Rng, m: &str, dims: [u64; 4], taps: u64, bias: Option<u64>, hot: bool) {
let n: u64 = dims.iter().product();
let a = [3i64, 255, 32767][r.below(3) as usize];
let w: Vec<i64> = (0..n).map(|_| r.signed(a)).collect();
// Most shifts 31; some others, from 24 to 40, with the multiplier scaled to match.
let shift: u64 = if r.below(4) == 0 { 24 + r.below(17) } else { 31 };
let root = (taps as f64).sqrt().max(1.0);
let mut mult = ((1u128 << shift) as f64 / (a as f64 * root)) as i64;
mult = (mult as f64 * [0.5, 1.0, 1.5, 2.0][r.below(4) as usize]) as i64;
if hot && r.below(64) == 0 {
mult = mult.saturating_mul(1 << 12);
}
self.0.insert(format!("{m}.weight"), json!({ "kind": "conv", "dims": dims, "mult": mult.max(1), "shift": shift, "w": w }));
if let Some(b) = bias {
self.q16(&format!("{m}.bias"), (0..b).map(|_| r.signed(1 << 17)).collect());
}
}
fn q16(&mut self, name: &str, v: Vec<i64>) {
self.0.insert(name.into(), json!({ "kind": "q16", "dims": [v.len()], "v": v }));
}
fn res(&mut self, r: &mut Rng, m: &str, c: u64, hot: bool) {
self.conv(r, &format!("{m}.conv1"), [c, c, 3, 3], 9 * c, Some(c), hot);
self.conv(r, &format!("{m}.conv2"), [c, c, 3, 3], 9 * c, Some(c), hot);
}
fn gate(&mut self, r: &mut Rng, m: &str, c: u64, hot: bool) {
self.conv(r, &format!("{m}.conv"), [c, c / 4, 3, 3], 9 * c / 4, None, hot);
self.conv(r, &format!("{m}.conv2"), [c, c, 1, 1], c, None, hot);
}
fn cab(&mut self, r: &mut Rng, c: u64, hot: bool) {
self.res(r, "CAB.residual_trunk.0", c, hot);
self.res(r, "CAB.residual_trunk.1", c, hot);
self.conv(r, "CAB.subscale", [c, c, 3, 3], 9 * c, Some(c), hot);
self.res(r, "CAB.residual_mask2.0", c, hot);
self.res(r, "CAB.residual_mask3.0", c, hot);
self.conv(r, "CAB.upscale", [c, c, 3, 3], 4 * c, Some(c), hot);
}
fn tab(&mut self, r: &mut Rng, m: &str, c: u64, hid: u64, hot: bool) {
let g: Vec<i64> = (0..c).map(|_| r.below(1 << 17) as i64).collect();
self.q16(&format!("{m}.prep_data.norm1.weight"), g);
self.q16(&format!("{m}.prep_data.norm1.bias"), (0..c).map(|_| r.signed(1 << 15)).collect());
self.conv(r, &format!("{m}.prep_data.conv1"), [3 * c, c, 1, 1], c, None, hot);
self.conv(r, &format!("{m}.prep_data.conv2"), [3 * c, 1, 3, 3], 9, None, hot);
self.q16(&format!("{m}.attn.temperature"), (0..4).map(|_| (1 << 16) + r.below(1 << 19) as i64).collect());
self.conv(r, &format!("{m}.attn.project_out"), [c, c, 1, 1], c, None, hot);
let g: Vec<i64> = (0..c).map(|_| r.signed(1 << 17)).collect();
self.q16(&format!("{m}.ffn.norm1.weight"), g);
self.q16(&format!("{m}.ffn.norm1.bias"), (0..c).map(|_| r.signed(1 << 15)).collect());
self.conv(r, &format!("{m}.ffn.project_in"), [2 * hid, c, 1, 1], c, None, hot);
self.conv(r, &format!("{m}.ffn.dwconv"), [2 * hid, 1, 3, 3], 9, None, hot);
self.conv(r, &format!("{m}.ffn.project_out"), [c, hid, 1, 1], hid, None, hot);
}
fn value(self) -> Value {
Value::Object(self.0)
}
}
/// A hyper decoder of internal width `w` over `cz` channels.
fn ptq_hyper(r: &mut Rng, cz: u64, w: u64, hot: bool) -> Value {
let mut f = NetFile::new();
f.conv(r, "conv1", [w, cz, 1, 1], cz, None, hot);
f.conv(r, "conv2", [w, w, 4, 4], 4 * w, Some(w), hot);
f.conv(r, "conv3", [w, w, 3, 3], 9 * w, Some(w), hot);
f.conv(r, "conv4", [4 * cz, w, 3, 3], 9 * w, Some(4 * cz), hot);
f.value()
}
/// Context stage `i` at the widths the direct stages read (160 throughout).
fn ptq_mcm(r: &mut Rng, i: u64, hot: bool) -> Value {
let mut f = NetFile::new();
if i > 0 {
f.conv(r, "conv.0", [160, 160 * i, 1, 1], 160 * i, None, hot);
f.conv(r, "conv.1", [160, 32, 3, 3], 288, Some(160), hot);
}
f.conv(r, "fusion_pred_net.conv1", [160, if i == 0 { 160 } else { 320 }, 1, 1], 320, None, hot);
f.conv(r, "fusion_pred_net.conv2", [160, 160, 1, 1], 160, None, hot);
f.conv(r, "fusion_pred_net.conv3", [160, 160, 1, 1], 160, None, hot);
f.value()
}
/// A synthesis network at the widths the direct networks read: luma 128 wide (its
/// feed-forward hidden width `hid`), chroma 64.
fn ptq_synth(r: &mut Rng, chroma: bool, hid: u64, hot: bool) -> Value {
let mut f = NetFile::new();
if !chroma {
let w = 128;
f.res(r, "first_stage.0.0", 160, hot);
f.conv(r, "first_stage.1", [160, w, 3, 3], 4 * 160, Some(w), hot);
f.gate(r, "first_stage.2", w, hot);
f.conv(r, "conv2_t", [w, w, 3, 3], 4 * w, Some(w), hot);
f.cab(r, w, hot);
f.gate(r, "iact2", w, hot);
f.conv(r, "conv3_t", [4 * w, w, 1, 1], w, Some(4 * w), hot);
f.conv(r, "TAM.ds_conv", [w, w, 3, 3], 9 * w, Some(w), hot);
f.tab(r, "TAM.TABs.0", w, hid, hot);
f.tab(r, "TAM.TABs.1", w, hid, hot);
f.conv(r, "TAM.us_conv", [w, w, 3, 3], 4 * w, Some(w), hot);
f.gate(r, "iact3", w, hot);
f.conv(r, "conv4_t", [w, 1, 3, 3], 4 * w, Some(1), hot);
} else {
let w = 64;
f.conv(r, "first_stage.conv1", [48, 256, 3, 3], 9 * 256, Some(48), hot);
f.conv(r, "conv2_t", [144, w, 3, 3], 4 * 144, Some(w), hot);
f.cab(r, w, hot);
f.gate(r, "iact2", w, hot);
f.conv(r, "conv3_t", [4 * w, w, 1, 1], w, Some(4 * w), hot);
f.tab(r, "TAM.TABs.0", w, hid, hot);
f.tab(r, "TAM.TABs.1", w, hid, hot);
f.gate(r, "iact3", w, hot);
f.conv(r, "conv4_t", [w, 8, 3, 3], 4 * w, Some(8), hot);
}
f.value()
}
/// The direct networks' weights, parsed as the decoder parses them.
fn ptq_weights(v: &Value) -> wai::jpegai_synthesis::Weights {
wai::jpegai_synthesis::Weights::from_json(v).unwrap()
}
/// Break one load check of §13.6.2 in a network file; the name of the check, or `None`
/// (one call in six).
fn ptq_break(r: &mut Rng, v: &mut Value) -> Option<&'static str> {
let o = v.as_object_mut().unwrap();
let convs: Vec<String> = o.iter().filter(|(_, e)| e["kind"] == "conv").map(|(k, _)| k.clone()).collect();
let k = convs[r.below(convs.len() as u64) as usize].clone();
Some(match r.below(6) {
0 => { o.remove(&k); "a member missing" }
1 => { o[&k]["shift"] = json!(128); "a shift of 128" }
2 => { o[&k]["w"].as_array_mut().unwrap().pop(); "a weight short" }
3 => { let d = o[&k]["dims"][0].as_u64().unwrap(); o[&k]["dims"][0] = json!(d + 1); "dims" }
4 => {
let b = k.replace(".weight", ".bias");
match o.get_mut(&b) {
Some(e) => { e["v"].as_array_mut().unwrap().push(json!(0)); "a bias long" }
None => return None,
}
}
_ => return None,
})
}
/// A random Q16 tensor of `n` values, of magnitude about `2^bits`.
fn q16_values(r: &mut Rng, n: usize, bits: u32) -> Vec<i64> {
(0..n).map(|_| r.signed(1 << bits)).collect()
}
/// The outcome of the networks through the operator set: equal to the direct networks',
/// or a refusal that is a planned break (checked here).
fn ptq_outcome(role: &str, new: Result<Vec<i64>, wai_ops::Refusal>, old: Option<Vec<i64>>, g: &wai_ops::Loaded, why: &dyn Fn() -> String) -> &'static str {
match (new, old) {
(Ok(a), Some(b)) => {
assert!(a == b, "{role}: the two decodes differ\n{}", why());
"alike"
}
(Err(r), _) => {
// A value past A32, at an input or at a checked tensor's producer.
assert_eq!(r.code, wai_ops::Code::FieldRange, "{role}: refused with {:?}\n{}", r, why());
let checked = match r.at {
wai_ops::Site::Input(k) => g.graph().tensors[g.graph().inputs[k as usize].tensor as usize].policy,
wai_ops::Site::Op(k) => g.graph().tensors[g.graph().ops[k as usize].outs[0] as usize].policy,
other => panic!("{role}: refused at {other:?}\n{}", why()),
};
assert_eq!(checked, wai_ops::graph::Policy::Check, "{role}: refused at a tensor that is not checked\n{}", why());
"past A32"
}
(Ok(_), None) => panic!("{role}: the direct network panicked where the operator set decodes\n{}", why()),
}
}
/// Random hyper decoders and context stages, with random hyper-latents and residuals:
/// `ŷ` through `rec_y` and `rec_uv` is the direct networks', or the operator set refuses
/// a value past `A32` (where the direct networks computed on, or panicked with overflow
/// checks on); a network file that fails a load check is refused, from its bytes and its
/// parsed form alike.
#[test]
fn jpegai_random_networks_reconstruct_alike() {
use wai::jpegai_ops::Plane;
use wai::jpegai_pipeline::{reconstruct_yhat, reconstruct_yhat_uv};
use wai::jpegai_synthesis::{decode_hyper, q16_from_symbols};
let mut outcomes = std::collections::BTreeMap::<&str, u64>::new();
let n = cases("jpegai_rec", 64, |r, seed| {
let hot = r.below(2) == 0;
let chroma = r.below(2) == 0;
let (plane, cz) = if chroma { (Plane::Chroma, 96u64) } else { (Plane::Luma, 160) };
let width = 1 + r.below(4);
let mut hyp = ptq_hyper(r, cz, width, hot);
let mut mcm: Vec<Value> = (0..4).map(|i| ptq_mcm(r, i, hot)).collect();
let broken = if r.below(8) == 0 {
if chroma || r.below(2) == 0 { ptq_break(r, &mut hyp) } else { let i = r.below(4) as usize; ptq_break(r, &mut mcm[i]) }
} else {
None
};
let why = || format!("seed {seed}: {plane:?}, broken {broken:?}");
let b = |v: &Value| serde_json::to_vec(v).unwrap();
// This plane's graph alone: `rec_y` from the hyper decoder and the four stages,
// or `rec_uv` from the hyper decoder.
let hyp_net = wai_ops::lower::jpegai_nets::parse_ptq(&b(&hyp));
let mcm_nets: Vec<_> = mcm.iter().map(|v| wai_ops::lower::jpegai_nets::parse_ptq(&b(v))).collect();
let graph = (|| -> Result<wai_ops::Graph, wai_ops::Refusal> {
let mut h = hyp_net.clone()?;
if chroma {
wai_ops::lower::jpegai_nets::rec(&mut h, None, true)
} else {
let mut m: [wai_ops::lower::jpegai_nets::PtqNet; 4] = [mcm_nets[0].clone()?, mcm_nets[1].clone()?, mcm_nets[2].clone()?, mcm_nets[3].clone()?];
wai_ops::lower::jpegai_nets::rec(&mut h, Some(&mut m), false)
}
})()
.and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64));
// The same, from the networks as this crate parses them, where its parser reads them.
let parsed = catch_unwind(AssertUnwindSafe(|| (ptq_weights(&hyp), mcm.iter().map(ptq_weights).collect::<Vec<_>>()))).ok();
if let Some((hw, mw)) = &parsed {
let g2 = (|| -> Result<wai_ops::Graph, wai_ops::Refusal> {
let mut h = hw.ptq()?;
if chroma {
wai_ops::lower::jpegai_nets::rec(&mut h, None, true)
} else {
let mut m = [mw[0].ptq()?, mw[1].ptq()?, mw[2].ptq()?, mw[3].ptq()?];
wai_ops::lower::jpegai_nets::rec(&mut h, Some(&mut m), false)
}
})()
.and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64));
assert_eq!(g2.as_ref().map(|g| g.graph().clone()).map_err(|e| e.code), graph.as_ref().map(|g| g.graph().clone()).map_err(|e| e.code), "the parsed networks lower as their bytes\n{}", why());
}
if let Some(k) = broken {
assert_eq!(graph.err().map(|e| e.code), Some(wai_ops::Code::ModelMismatch), "{k}: refused\n{}", why());
*outcomes.entry(k).or_default() += 1;
return;
}
let g = graph.unwrap_or_else(|e| panic!("networks under no planned break are refused: {e:?}\n{}", why()));
let (hw, mw) = parsed.expect("well-formed networks parse");
let c = cz as usize;
let (zh, zw) = ([1usize, 2, 4][r.below(3) as usize], [1usize, 2, 4][r.below(3) as usize]);
let z: Vec<i64> = (0..c * zh * zw).map(|_| r.below(64) as i64 - 31).collect();
let bits = [12u32, 20, 26][r.below(3) as usize];
let resi = q16_values(r, c * 16 * zh * zw, bits);
let new = wai_ops::run_with(&g, vec![wai_ops::Tensor::new(cz as u32, zh as u32, zw as u32, z.clone()), wai_ops::Tensor::new(cz as u32, 4 * zh as u32, 4 * zw as u32, resi.clone())], &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host)
.map(|mut o| o.swap_remove(0).data);
let old = catch_unwind(AssertUnwindSafe(|| {
let (psi, ..) = decode_hyper(&q16_from_symbols(&z), c, zh, zw, &hw);
if chroma {
reconstruct_yhat_uv(&resi, &psi, c, 4 * zh, 4 * zw)
} else {
reconstruct_yhat(&resi, &psi, c, 4 * zh, 4 * zw, [&mw[0], &mw[1], &mw[2], &mw[3]])
}
}))
.ok();
*outcomes.entry(ptq_outcome(if chroma { "rec_uv" } else { "rec_y" }, new, old, &g, &why)).or_default() += 1;
});
eprintln!("jpegai networks, rec: {n} cases; outcomes {outcomes:?}");
assert!(outcomes.get("alike").copied().unwrap_or(0) * 2 > n, "the generator must produce decodable cases: {outcomes:?}");
}
/// Random synthesis networks at the direct networks' widths, with random latents: the
/// planes through `syn_y` and `syn_uv` are the direct synthesis's, or the operator set
/// refuses a value past `A32`. These run at the whole picture's size, so the fixed count
/// is small; the nightly run gives them their time.
#[test]
fn jpegai_random_synthesis_networks_decode_alike() {
use wai::jpegai_synthesis::{decode_uv_f32, decode_y_f32};
let mut outcomes = std::collections::BTreeMap::<&str, u64>::new();
let n = cases("jpegai_syn", if cfg!(debug_assertions) { 2 } else { 8 }, |r, seed| {
let chroma = r.below(2) == 1;
let hot = r.below(2) == 0;
let hid = [32u64, 64, 256][r.below(3) as usize];
let v = ptq_synth(r, chroma, hid, hot);
let why = || format!("seed {seed}: chroma {chroma}, hid {hid}, hot {hot}");
let bytes = serde_json::to_vec(&v).unwrap();
let mut net = wai_ops::lower::jpegai_nets::parse_ptq(&bytes).unwrap();
let g = if chroma { wai_ops::lower::jpegai_nets::syn_uv(&mut net) } else { wai_ops::lower::jpegai_nets::syn_y(&mut net) }
.and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64))
.unwrap_or_else(|e| panic!("well-formed networks are refused: {e:?}\n{}", why()));
let w = ptq_weights(&v);
let bits = [12u32, 18, 22][r.below(3) as usize];
let ly = q16_values(r, 160 * 256, bits);
let luv = q16_values(r, 96 * 256, 18);
let mut inputs = vec![wai_ops::Tensor::new(160, 16, 16, ly.clone())];
if chroma {
inputs.push(wai_ops::Tensor::new(96, 16, 16, luv.clone()));
}
let new = wai_ops::run_with(&g, inputs, &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host).map(|mut o| o.swap_remove(0).data);
let old = catch_unwind(AssertUnwindSafe(|| {
let f = if chroma {
let mut cat = ly.clone();
cat.extend_from_slice(&luv);
decode_uv_f32(&cat, &w)
} else {
decode_y_f32(&ly, &w)
};
// Back to the Q16 integers the direct networks' denormalisation gave.
f.iter().map(|&x| (x as f64 * 65536.0) as i64).collect::<Vec<i64>>()
}))
.ok();
*outcomes.entry(ptq_outcome(if chroma { "syn_uv" } else { "syn_y" }, new, old, &g, &why)).or_default() += 1;
});
eprintln!("jpegai networks, synthesis: {n} cases; outcomes {outcomes:?}");
}

6
wai/wai-rs/tests/jpegai_cross_platform.rs
View file

@ -146,9 +146,9 @@ fn decode_matches_pinned_hashes() {
assert_eq!(hex(&Sha256::digest(&bytes)), want, "{} is not the pinned jpegai/v1 {name}", p.display());
}
let w = |n: &str| Weights::from_json(&jload(&dpath.join(n)));
let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n)));
let icci = IcciNet::from_json(&jload(&dpath.join("icci.json")));
let w = |n: &str| Weights::from_json(&jload(&dpath.join(n))).unwrap();
let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n))).unwrap();
let icci = IcciNet::from_json(&jload(&dpath.join("icci.json"))).unwrap();
let nv = jload(Path::new(&nf));
let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
let u8v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();

6
wai/wai-rs/tests/jpegai_decode_symbols.rs
View file

@ -81,9 +81,9 @@ fn link_the_decoder() {
black_box(jpegai_decode::check_scope as Check);
black_box(jpegai_decode::check_codestream as Check);
black_box(jpegai_decode::fp_env_is_default as fn() -> bool);
black_box(Weights::from_json as fn(&serde_json::Value) -> Weights);
black_box(IntNet::from_json as fn(&serde_json::Value) -> IntNet);
black_box(IcciNet::from_json as fn(&serde_json::Value) -> IcciNet);
black_box(Weights::from_json as fn(&serde_json::Value) -> Result<Weights, wai::jpegai_synthesis::BadMember>);
black_box(IntNet::from_json as fn(&serde_json::Value) -> Result<IntNet, wai::jpegai_synthesis::BadMember>);
black_box(IcciNet::from_json as fn(&serde_json::Value) -> Result<IcciNet, wai::jpegai_synthesis::BadMember>);
}
/// C99 `<math.h>` functions, besides the allowed ones. Each is also checked with the

2
wai/wai-rs/tests/jpegai_e2e_chain.rs
View file

@ -52,7 +52,7 @@ fn codestream_to_scale_indices() {
eprintln!("z from codestream: {} vals, {} non-zero", z.len(), z.iter().filter(|&&x| x != 0).count());
// 2. z → hyper-scale net → i_s
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap());
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()).unwrap();
let (i_s, co, ho, wo) = decode_hyper_scale(&z, c, zh, zw, &net);
eprintln!("i_s: [{co},{ho},{wo}] = {} vals", i_s.len());

4
wai/wai-rs/tests/jpegai_e2e_chroma.rs
View file

@ -48,7 +48,7 @@ fn env_path(k: &str) -> Option<String> {
std::env::var(k).ok().filter(|p| std::path::Path::new(p).exists())
}
fn weights(p: &str) -> Weights {
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap())
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()).unwrap()
}
fn snr(got_q16: &[i64], expected: &[f32]) -> f64 {
let (mut se, mut sig) = (0f64, 0f64);
@ -117,7 +117,7 @@ fn codestream_to_chroma_pixels() {
let z: Vec<i64> = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect();
// 2. z → hyper-scale → i_s, then the derived entropy inputs (shared with luma)
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&hs_uv).unwrap()).unwrap());
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&hs_uv).unwrap()).unwrap()).unwrap();
let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net);
assert_eq!(i_s.len(), channels * h * w, "i_s shape");

8
wai/wai-rs/tests/jpegai_e2e_frontend.rs
View file

@ -105,7 +105,7 @@ fn codestream_to_luma_pixels() {
let z: Vec<i64> = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect();
// 2. z → hyper-scale net → i_s
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap());
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()).unwrap();
let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net);
assert_eq!(i_s.len(), channels * h * w, "i_s shape");
@ -181,7 +181,7 @@ fn codestream_to_luma_pixels() {
// correctly lifted one scores ~70.
if let (Ok(hw), Some(psi_exp)) = (std::env::var("WAI_JPEGAI_HYPER_WEIGHTS"), v.get("psi_expected")) {
if std::path::Path::new(&hw).exists() {
let wt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&hw).unwrap()).unwrap());
let wt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&hw).unwrap()).unwrap()).unwrap();
let (psi, ph, pw) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &wt);
let exp: Vec<f32> = psi_exp.as_array().unwrap().iter()
.map(|x| x.as_f64().unwrap() as f32).collect();
@ -213,7 +213,7 @@ fn codestream_to_luma_pixels() {
if let Some(mcm_dir) = std::env::var_os("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
let load = |i: usize| -> Weights {
let p = std::path::Path::new(&mcm_dir).join(format!("mcm{i}_y.json"));
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap())
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()).unwrap()
};
let (m0, m1, m2, m3) = (load(0), load(1), load(2), load(3));
let resi_q16: Vec<i64> = resi_full.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect();
@ -261,7 +261,7 @@ fn codestream_to_luma_pixels() {
// deviation of at most 1 LSB, on at most 0.1% of samples.
if let (Ok(sw), Some(sy_exp)) = (std::env::var("WAI_JPEGAI_SYNTH_WEIGHTS"), v.get("synth_y_expected")) {
if std::path::Path::new(&sw).exists() {
let swt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&sw).unwrap()).unwrap());
let swt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&sw).unwrap()).unwrap()).unwrap();
let lat: Vec<i64> = y_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64).collect();
let y_plane = decode_y(&lat, &swt);
let exp: Vec<f32> = sy_exp.as_array().unwrap().iter()

8
wai/wai-rs/tests/jpegai_e2e_full.rs
View file

@ -48,10 +48,10 @@ fn jload(p: &std::path::Path) -> serde_json::Value {
serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
}
fn w_of(dir: &str, name: &str) -> Weights {
Weights::from_json(&jload(&std::path::Path::new(dir).join(name)))
Weights::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn int_of(dir: &str, name: &str) -> IntNet {
IntNet::from_json(&jload(&std::path::Path::new(dir).join(name)))
IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn u8v(v: &serde_json::Value, k: &str) -> Vec<u8> {
v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect()
@ -156,8 +156,8 @@ fn codestream_to_rgb_matches_reference_png() {
}
fn weights_from(get: &dyn Fn(&str) -> serde_json::Value) -> DecoderWeights {
let w = |n: &str| Weights::from_json(&get(n));
let i = |n: &str| IntNet::from_json(&get(n));
let w = |n: &str| Weights::from_json(&get(n)).unwrap();
let i = |n: &str| IntNet::from_json(&get(n)).unwrap();
let iw = get("icci.json");
DecoderWeights {
luma: PlaneWeights { hyper_scale: i("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },

6
wai/wai-rs/tests/jpegai_e2e_synthesis.rs
View file

@ -7,7 +7,7 @@
//! env: WAI_JPEGAI_E2E_CKPT = e2e_synth_checkpoint.json {yhat_y,yhat_uv,ypix,uvpix,y_lsbs}
//! WAI_JPEGAI_Y_WEIGHTS = synth16_y_weights.json
//! WAI_JPEGAI_UV_WEIGHTS = synth16_uv_weights.json
#![cfg(feature = "neural_int")]
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
use wai::jpegai_synthesis::{decode_y, decode_uv, Weights};
@ -38,8 +38,8 @@ fn synthesis_on_real_reference_latent() {
_ => { eprintln!("skip: set WAI_JPEGAI_Y_WEIGHTS + WAI_JPEGAI_UV_WEIGHTS"); return; }
};
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&ckpt).unwrap()).unwrap();
let wy = Weights::from_json(&serde_json::from_slice(&std::fs::read(&yw).unwrap()).unwrap());
let wu = Weights::from_json(&serde_json::from_slice(&std::fs::read(&uw).unwrap()).unwrap());
let wy = Weights::from_json(&serde_json::from_slice(&std::fs::read(&yw).unwrap()).unwrap()).unwrap();
let wu = Weights::from_json(&serde_json::from_slice(&std::fs::read(&uw).unwrap()).unwrap()).unwrap();
// Y: decode_y(yhat_y) vs ypix. Try both the pre-lsbs (yhat_y) and post-lsbs
// (y_lsbs) latent — whichever the synthesis actually consumes will hit ~58 dB.

4
wai/wai-rs/tests/jpegai_hyper_calibration.rs
View file

@ -6,7 +6,7 @@
//!
//! env: WAI_JPEGAI_HYPERDEC = path to common_modules/hyper_decoder.onnx,
//! WAI_JPEGAI_HYPERDEC_WEIGHTS = path to the exported 16-bit weights JSON.
#![cfg(feature = "neural")]
#![cfg(all(feature = "neural", feature = "legacy_int"))]
use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
@ -46,7 +46,7 @@ fn hyper_decoder_int_vs_float_snr() {
let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
let (pi, ph, pw) = decode_hyper(&li, c, h, w, &wt);
eprintln!("int hyper out = {} vals at [{},{}]", pi.len(), ph, pw);

2
wai/wai-rs/tests/jpegai_hyperscale_calibration.rs
View file

@ -36,7 +36,7 @@ fn hyper_scale_byte_exact_vs_reference() {
let zs: Vec<usize> = v["z_shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
let (c, h, w) = (zs[1], zs[2], zs[3]);
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap());
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()).unwrap();
let (mine, co, ho, wo) = decode_hyper_scale(&z, c, h, w, &net);
eprintln!("decode_hyper_scale(real z[{c},{h},{w}]) -> [{co},{ho},{wo}] = {} vals", mine.len());

4
wai/wai-rs/tests/jpegai_mcm_calibration.rs
View file

@ -6,7 +6,7 @@
//!
//! env: WAI_JPEGAI_MCM_DIR = dir with stage0..3.onnx,
//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
#![cfg(feature = "neural")]
#![cfg(all(feature = "neural", feature = "legacy_int"))]
use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
@ -56,7 +56,7 @@ fn mcm_stages_int_vs_float_snr() {
let (h, w) = (16usize, 16usize);
let load = |i: usize| -> Weights {
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap()).unwrap();
Weights::from_json(&v)
Weights::from_json(&v).unwrap()
};
let mut worst = f64::INFINITY;

31
wai/wai-rs/tests/jpegai_memory.rs
View file

@ -1,15 +1,16 @@
//! What `wai.image.jpegai`'s exact decoder holds now that its scale path and entropy
//! decodes run through the integer operator set, against the direct decode it replaced
//! What `wai.image.jpegai`'s exact decoder holds now that every integer stage runs
//! through the integer operator set, against the direct decode it replaced
//! (`jpegai_decode::legacy`, feature `legacy_int`), measured with a global allocator that
//! tracks the bytes live and their peak.
//!
//! The capability declares no cost. So what the decode holds is held to the direct
//! decode: with the bundle's scale path loaded beforehand (`jpegai_decode::load`, once per
//! bundle, as the sink's store of loaded sets keeps it), a decode through the operator set
//! holds no more than the direct decode of the same codestream, on the reference kernels;
//! and so does each plane's front half (its integers through §13.6.1), measured alone.
//! The scale path itself (two graphs and the latent tables) is the set's, the same for
//! every codestream, and is reported. The inputs are the `jpegai/v1` bundle and the test
//! decode: with the bundle's scale path and networks loaded beforehand
//! (`jpegai_decode::load` and `jpegai_decode::load_networks`, once per bundle, as the
//! sink's store of loaded sets keeps them), a decode through the operator set holds no
//! more than the direct decode of the same codestream, on the reference kernels; and so
//! does each plane's front half (its integers through §13.6.1), measured alone. The
//! scale path (two graphs and the latent tables) and the networks (four graphs, §13.6.2)
//! are the set's, the same for every codestream, and are reported. The inputs are the `jpegai/v1` bundle and the test
//! codestream (`WAI_JPEGAI_CODESTREAM`, `WAI_JPEGAI_WEIGHTS_DIR`, `WAI_JPEGAI_NORM`), so
//! the test runs where they are fetched (`wai-jpegai-e2e.yml`) and skips loudly elsewhere.
//!
@ -81,13 +82,13 @@ fn the_decode_holds_no_more_than_the_direct_decode() {
return;
};
let d = PathBuf::from(dir);
let w = |n: &str| Weights::from_json(&json(&d.join(n)));
let net = |n: &str| IntNet::from_json(&json(&d.join(n)));
let w = |n: &str| Weights::from_json(&json(&d.join(n))).unwrap();
let net = |n: &str| IntNet::from_json(&json(&d.join(n))).unwrap();
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
icci: wai::jpegai_icci::IcciNet::from_json(&json(&d.join("icci.json"))),
icci: wai::jpegai_icci::IcciNet::from_json(&json(&d.join("icci.json"))).unwrap(),
};
let nv = json(Path::new(&nf));
let ints = |k: &str| -> Vec<i64> { nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() };
@ -111,13 +112,15 @@ fn the_decode_holds_no_more_than_the_direct_decode() {
let (direct, direct_peak) = peak_during(|| legacy::decode(&data, &weights, &norm).expect("the direct decode"));
drop(direct);
let (sp, set) = peak_during(|| jpegai_decode::load(&weights, &norm).expect("the scale path loads"));
let (img, held) = peak_during(|| jpegai_decode::decode_loaded(&data, &weights, &norm, &sp, &mut |_, _| {}).expect("the decode"));
let (nets, nets_load) = peak_during(|| jpegai_decode::load_networks(&weights).expect("the networks load"));
let (img, held) = peak_during(|| jpegai_decode::decode_loaded(&data, &weights.icci, &norm, &sp, &nets, &mut |_, _| {}).expect("the decode"));
drop(img);
let (img, whole) = peak_during(|| jpegai_decode::decode(&data, &weights, &norm).expect("the decode"));
drop(img);
eprintln!(
"jpegai/v1, reference kernels: the direct decode held {direct_peak} bytes; with the scale path loaded beforehand ({set} bytes to load, {} held), the decode through the operator set held {held}; loading and decoding in one call held {whole}",
sp.heap_bytes()
"jpegai/v1, reference kernels: the direct decode held {direct_peak} bytes; with the scale path ({set} bytes to load, {} held) and the networks ({nets_load} bytes to load, {} held) loaded beforehand, the decode through the operator set held {held}; loading and decoding in one call held {whole}",
sp.heap_bytes(),
nets.heap_bytes()
);
assert!(held <= direct_peak, "the decode through the operator set held {held} bytes; the direct decode held {direct_peak}");

608
wai/wai-rs/tests/jpegai_nets_ops.rs Normal file
View file

@ -0,0 +1,608 @@
//! `wai.image.jpegai`'s other integer networks through the integer operator set: the
//! hyper decoder, the four context stages and the synthesis, as `jpegai-int/1` lowers
//! them (`extensions/integer-operators.md` §13.6.2) to the graphs `rec_y`, `rec_uv`,
//! `syn_y` and `syn_uv`, and as this crate loads them (`jpegai_ops::Networks`).
//!
//! The committed bundle (`ops-conformance/sets/jpegai`, synthetic, at small widths) has
//! its lowered graphs and certificate under `ops-conformance/certificates/`, where
//! `wai-verify` and the independent checker (`tools/wai_ops_check.py`) reproduce them
//! byte for byte. The `jpegai/v1` bundle of the models root is not committed; its
//! certificate is, by digest, and the tests that read it take its files and the test
//! codestream from `WAI_JPEGAI_WEIGHTS_DIR`, `WAI_JPEGAI_NORM` and
//! `WAI_JPEGAI_CODESTREAM` (`tools/wai_jpegai_fetch_e2e_inputs.sh`), with
//! `WAI_REQUIRE_SOTA_VECTORS` deciding whether their absence is a skip or a failure.
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
#[path = "support/require_vectors.rs"]
mod require_vectors;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use serde_json::Value;
use wai::jpegai_decode::{DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_ops::{self, Networks, Plane};
use wai::jpegai_synthesis::{IntNet, Weights};
use wai_ops::lower::jpegai_nets::NET_MEMBERS;
use wai_ops::{Profile, Tensor};
const CAP: &str = "wai.image.jpegai";
const SYNTHETIC: &str = "1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9";
const V1: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37";
const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)";
const ROLES: [&str; 4] = ["rec_y", "rec_uv", "syn_y", "syn_uv"];
fn wai() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(p: &Path) -> Vec<u8> {
std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
}
fn json(b: &[u8]) -> Value {
serde_json::from_slice(b).unwrap()
}
fn sha256(b: &[u8]) -> String {
wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect()
}
fn i64s(v: &Value) -> Vec<i64> {
v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}
/// A bundle's members as bytes, by their names in the set (`norm.json`, `weights/…`).
struct Members(BTreeMap<String, Vec<u8>>);
impl Members {
fn synthetic() -> Members {
let d = wai().join("ops-conformance/sets/jpegai");
let mut m = BTreeMap::new();
for n in ["norm.json", "weights/hs_y.json", "weights/hs_uv.json"].into_iter().chain(NET_MEMBERS) {
m.insert(n.to_string(), read(&d.join(n)));
}
Members(m)
}
/// The `jpegai/v1` bundle: its twelve members, from the variables the whole-decode
/// tests read (a directory of `<name>.json`, and `norm.json`).
fn pinned(test: &str) -> Option<Members> {
let [dir, norm] = require_vectors::env_paths(test, ["WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
let d = PathBuf::from(dir);
let mut m = BTreeMap::new();
m.insert("norm.json".to_string(), read(Path::new(&norm)));
for n in ["hs_y", "hs_uv", "hyp_y", "hyp_uv", "synth_y", "synth_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "icci"] {
m.insert(format!("weights/{n}.json"), read(&d.join(format!("{n}.json"))));
}
Some(Members(m))
}
fn get(&self, n: &str) -> &[u8] {
&self.0[n]
}
fn bytes(&self) -> wai_ops::lower::JpegaiBytes<'_> {
wai_ops::lower::JpegaiBytes::from_members(|n| self.0.get(n).map(Vec::as_slice)).expect("every member the lowering reads")
}
fn nets(&self) -> [&[u8]; 8] {
NET_MEMBERS.map(|n| self.get(n))
}
fn digest(&self) -> String {
let m: BTreeMap<String, &[u8]> = self.0.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
wai::int_set_cache::SetKey::listing_digest(&m).unwrap().iter().map(|x| format!("{x:02x}")).collect()
}
/// The networks as this crate parses them (`Weights::from_json`).
fn weights(&self, n: &str) -> Weights {
Weights::from_json(&json(self.get(n))).unwrap()
}
}
/// Gate 1, on one bundle: the lowering of its bytes writes `hs_y`, `hs_uv`, then the four
/// network graphs; the networks as this crate parses them, converted
/// (`jpegai_ops::load_networks`), are the graphs `wai-ops` lowers from the bytes; the
/// committed certificate `digest` is those six graphs', with the verdict `acc64`, and
/// `acc32` for the scale path alone (§11's `role_profiles`).
fn lowers_to_its_certificate(name: &str, m: &Members, digest: &str) {
let lowered = wai_ops::lower::jpegai(&m.bytes()).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
assert_eq!(roles, ["hs_y", "hs_uv", "rec_y", "rec_uv", "syn_y", "syn_uv"], "{name}: the roles, in §13.0's order");
let w = |n: &str| m.weights(n);
let dw = DecoderWeights {
luma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_y.json"))).unwrap(), hyper_decoder: w("weights/hyp_y.json"), synthesis: w("weights/synth_y.json") },
chroma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_uv.json"))).unwrap(), hyper_decoder: w("weights/hyp_uv.json"), synthesis: w("weights/synth_uv.json") },
mcm_y: [w("weights/mcm0_y.json"), w("weights/mcm1_y.json"), w("weights/mcm2_y.json"), w("weights/mcm3_y.json")],
icci: wai::jpegai_icci::IcciNet { luma: empty_path(), chroma: empty_path(), conv_last_v: empty_conv(), scale_v: 0.0 },
};
let parsed = jpegai_ops::nets_of(&dw).unwrap_or_else(|r| panic!("{name}: does not convert: {r:?}"));
assert!(parsed == wai_ops::lower::jpegai_nets::parse_nets(m.nets()).unwrap(), "{name}: the converted networks are the ones wai-ops parses from the bytes");
let nets = jpegai_ops::load_networks(&dw).unwrap_or_else(|r| panic!("{name}: refused under acc64: {r:?}"));
let from_bytes = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
let cert = json(&read(&wai().join(format!("ops-conformance/certificates/{CAP}/{digest}.json"))));
assert_eq!(cert["set_sha256"].as_str(), Some(digest), "{name}");
for (i, role) in ROLES.iter().enumerate() {
let g = lowered.get(role).unwrap();
assert!(*nets.graph(role).unwrap().graph() == *g, "{name}: {role} as loaded from the parsed networks");
assert!(*from_bytes.graph(role).unwrap().graph() == *g, "{name}: {role} as loaded from the bytes");
assert!(!nets.graph(role).unwrap().meets(Profile::Acc32), "{name}: {role} meets acc64 and not acc32");
let c = &cert["graphs"][i + 2];
assert_eq!(c["name"], *role, "{name}");
assert_eq!(c["graph_sha256"].as_str(), Some(sha256(&g.to_bytes()).as_str()), "{name}: the certificate is {role}'s");
let lo = (i32::MIN as i64).to_string();
let hi = (i32::MAX as i64).to_string();
for x in c["inputs"].as_array().unwrap() {
let range = (x["lo"].as_str().unwrap(), x["hi"].as_str().unwrap());
assert!(x["policy"] == "check" && (range == ("-31", "32") || range == (lo.as_str(), hi.as_str())), "{name}: {role}'s input ranges: {range:?}");
}
}
for (i, role) in ["hs_y", "hs_uv"].iter().enumerate() {
let c = &cert["graphs"][i];
assert_eq!(c["name"], *role);
assert!(c["nodes"].as_array().unwrap().iter().all(|n| n["fits"].as_array().unwrap().iter().any(|p| p == "acc32")), "{name}: every node of {role} fits acc32");
}
let v = &cert["verdict"];
assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the set meets acc64 and not acc32");
assert_eq!(v["acc32_first_failure"]["graph"], "rec_y", "{name}: the first acc32 failure lies past the scale path");
}
fn empty_conv() -> wai::jpegai_icci::Conv {
wai::jpegai_icci::Conv { w: vec![], b: vec![], cin: 0, cout: 0, kh: 0, kw: 0, ph: 0, pw: 0 }
}
fn empty_path() -> wai::jpegai_icci::IcciPath {
wai::jpegai_icci::IcciPath { conv_first: empty_conv(), blocks: vec![], conv_last: empty_conv(), scale: 0.0 }
}
/// The committed synthetic bundle: its four network graphs are the committed ones, and
/// its set digest is the listing of its eleven members.
#[test]
fn the_committed_bundle_lowers_to_its_committed_networks() {
let m = Members::synthetic();
assert_eq!(m.digest(), SYNTHETIC, "the set's digest");
let dir = wai().join(format!("ops-conformance/certificates/{CAP}/{SYNTHETIC}"));
let lowered = wai_ops::lower::jpegai(&m.bytes()).unwrap();
for role in ROLES {
assert!(lowered.get(role).unwrap().to_bytes() == read(&dir.join(format!("{role}.whg1"))), "the lowering of the bytes differs from the committed {role}.whg1");
}
lowers_to_its_certificate("synthetic", &m, SYNTHETIC);
}
/// The `jpegai/v1` bundle: its networks lower to its committed certificate's graphs.
#[test]
fn the_pinned_bundle_lowers_its_networks_to_its_certificate() {
let Some(m) = Members::pinned("jpegai_nets_ops::the_pinned_bundle_lowers_its_networks_to_its_certificate") else {
return;
};
assert_eq!(m.digest(), V1, "the set's digest");
lowers_to_its_certificate("jpegai/v1", &m, V1);
eprintln!("jpegai/v1: rec_y, rec_uv, syn_y and syn_uv lower to the committed certificate {V1} and meet acc64");
}
/// The whole `jpegai/v1` bundle and the test codestream, as the decoders parse them.
struct Whole {
data: Vec<u8>,
weights: DecoderWeights,
norm: NormTables,
}
impl Whole {
fn pinned(test: &str) -> Option<Whole> {
let [cs, dir, nf] = require_vectors::env_paths(test, ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
let d = PathBuf::from(dir);
let j = |n: &str| json(&read(&d.join(n)));
let w = |n: &str| Weights::from_json(&j(n)).unwrap();
let net = |n: &str| IntNet::from_json(&j(n)).unwrap();
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")).unwrap(),
};
let nv = json(&read(Path::new(&nf)));
let three = |k: &str| -> [f64; 3] { std::array::from_fn(|i| nv[k][i].as_f64().unwrap()) };
let norm = NormTables {
pmf: nv["pmf"].as_array().unwrap().iter().map(i64s).collect(),
bounds: i64s(&nv["bounds"]),
per_channel: nv["per_channel"].as_u64().unwrap() as usize,
cdfs_y: i64s(&nv["cdfs_y"]),
cdfs_uv: i64s(&nv["cdfs_uv"]),
gain_y: i64s(&nv["gain_y"]),
gain_uv: i64s(&nv["gain_uv"]),
lsbs_t0: i64s(&nv["lsbs_t0"]).iter().map(|&x| x as i32).collect(),
lsbs_t1: i64s(&nv["lsbs_t1"]).iter().map(|&x| x as i32).collect(),
lef_thr: three("lef_thr0").map(|x| x as i64),
lef_mag: three("lef_mag0").map(|x| x as f32),
};
Some(Whole { data: read(Path::new(&cs)), weights, norm })
}
}
/// What the networks of §13.6.2 read from the test codestream, and what the direct
/// decode makes of it: for each plane `z`, the dequantised residual in Q16 (the input
/// the float stages hand `rec_*`), the direct decode's `ŷ`, and the LSBS-refined latent
/// in Q16 (the input of `syn_*`); then the direct synthesis's planes, as the float
/// values the decode reports.
struct Inputs {
z: [Vec<i64>; 2],
resi: [Vec<i64>; 2],
yhat: [Vec<i64>; 2],
latent: [Vec<i64>; 2],
syn_y: Vec<f32>,
syn_uv: Vec<f32>,
}
fn q16(f: &[f32]) -> Vec<i64> {
f.iter().map(|&x| wai::float_repro::round_half_away(x as f64 * 65536.0) as i64).collect()
}
/// The decode's own steps up to each network (`jpegai_decode`'s, restated), with the
/// direct networks (`jpegai_synthesis`, `jpegai_pipeline`).
fn inputs(wh: &Whole) -> Inputs {
use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat, reconstruct_yhat_uv};
use wai::jpegai_synthesis::{decode_hyper, decode_uv_f32, decode_y_f32, q16_from_symbols};
let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
let segs = parse_container(&wh.data);
let pay = |m| {
let s = find_segment(&segs, m).unwrap();
wh.data[s.start..s.start + s.len].to_vec()
};
let ph = parse_picture_header(&wh.data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
let mut out: Vec<(Vec<i64>, Vec<i64>, Vec<i64>, Vec<i64>)> = vec![];
for (plane, lat_m, pw, gain) in [(Plane::Luma, Marker::Sorp, &wh.weights.luma, &wh.norm.gain_y), (Plane::Chroma, Marker::Sors, &wh.weights.chroma, &wh.norm.gain_uv)] {
let c = plane.channels();
let mut grfs = ph.components[plane as usize].grfs_channel_flags.clone();
grfs.resize(c, false);
let ints = jpegai_ops::plane(&sp, plane, &pay(Marker::Soz), &pay(lat_m), &grfs).unwrap();
let y_q: Vec<i16> = ints.y_q.iter().map(|&q| q as i16).collect();
let resi = dequantize_resi(&y_q, &ints.likely, &grfs, &gain_unit_scaler_vec(gain, 0), c);
let resi_q = q16(&resi);
let (psi, ..) = decode_hyper(&q16_from_symbols(&ints.z), c, 4, 4, &pw.hyper_decoder);
let yhat = if plane == Plane::Luma {
let m = &wh.weights.mcm_y;
reconstruct_yhat(&resi_q, &psi, c, 16, 16, [&m[0], &m[1], &m[2], &m[3]])
} else {
reconstruct_yhat_uv(&resi_q, &psi, c, 16, 16)
};
let yf: Vec<f32> = yhat.iter().map(|&v| v as f64 as f32 / 65536.0).collect();
let lat = lsbs_refine(&yf, &resi, &ints.likely, &wh.norm.lsbs_t0, &wh.norm.lsbs_t1, c, 16, 16);
out.push((ints.z, resi_q, yhat, q16(&lat)));
}
let [(zy, ry, yy, ly), (zuv, ruv, yuv, luv)]: [_; 2] = out.try_into().unwrap();
let syn_y = decode_y_f32(&ly, &wh.weights.luma.synthesis);
let mut cat = ly.clone();
cat.extend_from_slice(&luv);
let syn_uv = decode_uv_f32(&cat, &wh.weights.chroma.synthesis);
Inputs { z: [zy, zuv], resi: [ry, ruv], yhat: [yy, yuv], latent: [ly, luv], syn_y, syn_uv }
}
fn floats(q: &[i64]) -> Vec<f32> {
q.iter().map(|&v| v as f32 / 65536.0).collect()
}
/// Gate 2, the networks: on the test codestream, each plane's `ŷ` through `rec_y` and
/// `rec_uv` is the direct hyper decoder's and context stages', every value, and the
/// planes through `syn_y` and `syn_uv` are the direct synthesis's, bit for bit as the
/// float values the decode reports.
#[test]
fn the_networks_decode_through_the_lowering_as_directly() {
let Some(wh) = Whole::pinned("jpegai_nets_ops::the_networks_decode_through_the_lowering_as_directly") else {
return;
};
let i = inputs(&wh);
let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
for (k, plane) in [Plane::Luma, Plane::Chroma].into_iter().enumerate() {
let yhat = jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), i.resi[k].clone()).unwrap();
assert!(yhat == i.yhat[k], "{plane:?}: ŷ through the lowering differs from the direct hyper decoder and context stages");
eprintln!("{plane:?}: ŷ ({} values) equals the direct decode's", yhat.len());
}
let y = floats(&jpegai_ops::synthesize_y(&nets, i.latent[0].clone()).unwrap());
assert!(y.iter().map(|v| v.to_bits()).eq(i.syn_y.iter().map(|v| v.to_bits())), "Y through syn_y differs from the direct synthesis");
let uv = floats(&jpegai_ops::synthesize_uv(&nets, i.latent[0].clone(), i.latent[1].clone()).unwrap());
assert!(uv.iter().map(|v| v.to_bits()).eq(i.syn_uv.iter().map(|v| v.to_bits())), "U, V through syn_uv differ from the direct synthesis");
eprintln!("Y ({}) and U, V ({}) through the synthesis graphs equal the direct synthesis's", y.len(), uv.len());
}
/// The pinned decode's stage hashes (`tests/jpegai_cross_platform.rs`' `PINNED`): the
/// float values after each stage, and the RGB bytes, of the `jpegai/v1` decode.
const PINNED: [(&str, &str); 8] = [
("latent_y", "5889f3806308496c64850aee04f3706178f92d1e0edcda1e6975fe256779061e"),
("latent_uv", "be89a76f726ef362f267aba3dee8f61594cc512576a3864f52f665cc5b6ab77b"),
("synthesis", "d56695726bb6ebed403199f60a79a5698734eef1e5ca2813859d490c0882d2a4"),
("efl", "1783ea77bd4f6b0ac6a599137577a18db2a86753aa3eb3d79ca987e1cf25792e"),
("icci", "3566c51c4042a749d7826df812663c4e3ff7a9e807cbd0c3bb0545c9dbd79f2b"),
("efn", "09dc833c7a6b91ce08da20adb8bbc55e2eeedba7d8a2dc4b133ab0915b267de7"),
("lef", "e2b20c72ac6a59dd12067f75562c82715061f6d7a3444ca075fc0b11e3f56468"),
("rgb", "dbc494a5a670914edf55c07c48afc0dd348cd05d3a82fff293bf384b2bb69ea3"),
];
/// Gate 2, the decode: the whole decode with every integer stage through the operator
/// set (the scale path, the entropy decodes and the four network graphs) gives the
/// direct decode's float values after every stage, bit for bit, and its RGB bytes; and
/// those are the pinned decode's: the `Synthesis` stage hash and the RGB SHA-256 are
/// unchanged.
#[test]
fn the_whole_decode_through_the_networks_is_the_pinned_decode() {
let Some(wh) = Whole::pinned("jpegai_nets_ops::the_whole_decode_through_the_networks_is_the_pinned_decode") else {
return;
};
let bits = |planes: &[&[f32]]| -> String { sha256(&planes.iter().flat_map(|p| p.iter().flat_map(|v| v.to_bits().to_le_bytes())).collect::<Vec<u8>>()) };
let mut direct = vec![];
let a = wai::jpegai_decode::legacy::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| direct.push((s.name(), bits(p)))).expect("the direct decode");
direct.push(("rgb", sha256(&a.data)));
let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
let mut ops = vec![];
let b = wai::jpegai_decode::decode_loaded(&wh.data, &wh.weights.icci, &wh.norm, &sp, &nets, &mut |s, p| ops.push((s.name(), bits(p)))).expect("the decode through the operator set");
// The decode a sink makes, which loads the bundle itself, is the same.
let c = wai::jpegai_decode::decode(&wh.data, &wh.weights, &wh.norm).expect("the decode");
assert!(c.data == b.data, "jpegai_decode::decode");
ops.push(("rgb", sha256(&b.data)));
assert_eq!(ops, direct, "the stages through the networks' graphs");
let pinned: Vec<(&str, String)> = PINNED.iter().map(|(n, h)| (*n, h.to_string())).collect();
assert_eq!(ops, pinned, "the decode through the networks' graphs is not the pinned decode");
eprintln!("every stage hash ({}) through the networks' graphs is the pinned decode's, the synthesis {} and the RGB {} among them", ops.len(), PINNED[2].1, PINNED[7].1);
}
/// The accumulator class, measured (§13.6.2's table): the four graphs run on the test
/// codestream's inputs through the interpreter alone, with every value each operator
/// forms held to its certificate's interval (`wai_ops::bounds::Within`), and the largest
/// magnitude each kind of sum reaches, per graph, printed beside the certificate's bound.
/// Nothing a decode forms leaves its interval, and no sum passes `2^63`.
#[test]
fn the_networks_keep_to_their_certificate_on_the_test_codestream() {
let Some(wh) = Whole::pinned("jpegai_nets_ops::the_networks_keep_to_their_certificate_on_the_test_codestream") else {
return;
};
let i = inputs(&wh);
let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
let runs: [(&str, Vec<Tensor>); 4] = [
("rec_y", vec![Tensor::new(160, 4, 4, i.z[0].clone()), Tensor::new(160, 16, 16, i.resi[0].clone())]),
("rec_uv", vec![Tensor::new(96, 4, 4, i.z[1].clone()), Tensor::new(96, 16, 16, i.resi[1].clone())]),
("syn_y", vec![Tensor::new(160, 16, 16, i.latent[0].clone())]),
("syn_uv", vec![Tensor::new(160, 16, 16, i.latent[0].clone()), Tensor::new(96, 16, 16, i.latent[1].clone())]),
];
const SUMS: [&str; 9] = ["conv.acc", "convT.acc", "rq.t", "rq.p", "mul.p", "gram.acc", "apply.acc", "ln.T", "l2.T"];
for (role, x) in runs {
let g = nets.graph(role).unwrap();
let sizes = wai_ops::plan::load_plan(g.graph()).unwrap();
let b = wai_ops::bounds::pass(g.graph(), sizes.as_deref()).unwrap();
let mut w = wai_ops::bounds::Within::new(&b);
wai_ops::interp::run_observed(g, &x, &[], &mut wai_ops::Sequence::new(), &mut w).unwrap_or_else(|r| panic!("{role}: {r:?}"));
assert!(w.outside.is_empty(), "{role}: values outside the certificate: {:?}", w.outside);
assert!(w.unrecorded.is_empty(), "{role}: keys the certificate does not record: {:?}", w.unrecorded);
// The largest magnitude seen under each sum, over the graph's nodes, with the
// certificate's bound for the same key over the same nodes.
for key in SUMS {
let mut seen: Option<(u128, u32)> = None;
let mut bound: u128 = 0;
for (k, n) in b.nodes.iter().enumerate() {
if let Some((lo, hi)) = w.extreme(k as u32, key) {
let m = lo.unsigned_abs().max(hi.unsigned_abs());
if seen.is_none_or(|(s, _)| m > s) {
seen = Some((m, k as u32));
}
if let Some(iv) = n.inter.iter().find(|x| x.key == key) {
bound = bound.max(iv.iv.bits() as u128);
}
}
}
if let Some((m, at)) = seen {
let bits = 128 - m.leading_zeros();
assert!(m < 1u128 << 63 || key.starts_with("rq.p") || key.starts_with("ln.") || key.starts_with("l2."), "{role}: {key} reaches {m} at node {at}");
eprintln!("{role:7} {key:10} largest |value| {m:>22} ({bits:>2} bits) at node {at:>3}; certified within {bound} bits");
}
}
// Every activation the graph writes (each producer result), and its inputs.
let mut act: u128 = 0;
for k in 0..b.nodes.len() as u32 {
if let Some((lo, hi)) = w.extreme(k, "out.0") {
act = act.max(lo.unsigned_abs()).max(hi.unsigned_abs());
}
}
let inp = x.iter().flat_map(|t| t.data.iter()).map(|v| v.unsigned_abs() as u128).max().unwrap_or(0);
assert!(act < 1 << 31 && inp < 1 << 31, "{role}: an activation or an input reaches 2^31");
eprintln!("{role:7} largest |input| {inp} ({} bits), largest |activation| {act} ({} bits); every one within i32", 128 - inp.leading_zeros(), 128 - act.leading_zeros());
eprintln!("{role:7} {} values held to the certificate", w.seen);
}
}
/// Gate 5, the fast kernels: under every kernel policy, each of the four graphs gives the
/// interpreter's values, on the committed synthetic networks with inputs over their
/// ranges at small sizes (in the fast-kernel steps) and, where the bundle is supplied, on
/// the test codestream's inputs. (`jpegai_ops`' own test counts the convolutions that
/// ran on the fast kernel.)
#[test]
fn every_kernel_policy_runs_the_networks_alike() {
use wai::int_kernels::{with_policy, Policy};
const POLICIES: [Policy; 5] = [Policy::reference(), Policy::fast(1), Policy::fast(2), Policy::fast(5), Policy::fast(0)];
let m = Members::synthetic();
let syn = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
let mut s: u64 = 0x0D2C_007B;
let mut range = |n: usize, lo: i64, hi: i64| -> Vec<i64> {
(0..n).map(|_| {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
lo + ((s >> 11) % (hi - lo + 1) as u64) as i64
}).collect()
};
let mut cases: Vec<(&Networks, &str, Vec<Tensor>)> = vec![];
for k in 0..3 {
let a = 1i64 << (16 + 4 * k);
cases.push((&syn, "rec_y", vec![Tensor::new(160, 1, 1, range(160, -31, 32)), Tensor::new(160, 4, 4, range(160 * 16, -a, a))]));
cases.push((&syn, "rec_uv", vec![Tensor::new(96, 2, 1, range(192, -31, 32)), Tensor::new(96, 8, 4, range(96 * 32, -a, a))]));
cases.push((&syn, "syn_y", vec![Tensor::new(160, 2, 2, range(640, -a, a))]));
cases.push((&syn, "syn_uv", vec![Tensor::new(160, 2, 2, range(640, -a, a)), Tensor::new(96, 2, 2, range(384, -a, a))]));
}
let pinned = Whole::pinned("jpegai_nets_ops::every_kernel_policy_runs_the_networks_alike");
let real = pinned.as_ref().map(|wh| (inputs(wh), jpegai_ops::load_networks(&wh.weights).unwrap()));
if let Some((i, n)) = &real {
cases.push((n, "rec_y", vec![Tensor::new(160, 4, 4, i.z[0].clone()), Tensor::new(160, 16, 16, i.resi[0].clone())]));
cases.push((n, "rec_uv", vec![Tensor::new(96, 4, 4, i.z[1].clone()), Tensor::new(96, 16, 16, i.resi[1].clone())]));
cases.push((n, "syn_y", vec![Tensor::new(160, 16, 16, i.latent[0].clone())]));
cases.push((n, "syn_uv", vec![Tensor::new(160, 16, 16, i.latent[0].clone()), Tensor::new(96, 16, 16, i.latent[1].clone())]));
}
let (mut ran, mut refused) = (0, 0);
for (n, role, x) in &cases {
let g = n.graph(role).unwrap();
let run = || wai_ops::run_with(g, x.clone(), &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host).map(|mut o| o.swap_remove(0).data);
let want = with_policy(POLICIES[0], run);
for p in &POLICIES[1..] {
assert!(with_policy(*p, run) == want, "{role} {p:?}: the outputs differ from the interpreter alone");
}
ran += 1;
refused += want.is_err() as usize;
}
eprintln!("{ran} runs of the four graphs alike under every policy ({refused} refused alike){}", if real.is_some() { ", the test codestream's among them" } else { "" });
}
/// One network file of the committed synthetic bundle changed by `edit` (the file's JSON
/// object, its members by name), as bytes; and whether this crate's parser
/// (`Weights::from_json`) reads it, which it does only where every member of kind `conv`
/// or `q16` has its fields of their types (and otherwise refuses it).
fn edited_net(file: &str, edit: impl Fn(&mut serde_json::Map<String, Value>)) -> (Vec<u8>, bool) {
let m = Members::synthetic();
let mut v = json(m.get(file));
edit(v.as_object_mut().unwrap());
let typed = v.as_object().unwrap().values().all(|e| match e.get("kind").and_then(Value::as_str) {
Some("conv") => e["w"].as_array().is_some_and(|a| a.iter().all(Value::is_i64)) && e["dims"].as_array().is_some_and(|a| a.iter().all(Value::is_u64)) && e["mult"].is_i64() && e["shift"].is_u64(),
Some("q16") => e["v"].as_array().is_some_and(|a| a.iter().all(Value::is_i64)),
_ => true,
});
(serde_json::to_vec(&v).unwrap(), typed)
}
/// The eight network files of the committed bundle with `file` replaced by `b`.
fn nets_with<'a>(m: &'a Members, file: &str, b: &'a [u8]) -> [&'a [u8]; 8] {
NET_MEMBERS.map(|n| if n == file { b } else { m.get(n) })
}
/// Gate 3, the bundle: each load check of §13.6.2, broken one change at a time in the
/// committed bundle's network files, is refused `model_mismatch` by the load from the
/// files' bytes, and by this crate's parser or the load from the networks as it parses
/// them. The direct decoder made none of these checks (integer-operators §15.3 item 4):
/// its loaders panicked on a mistyped member, and it computed with whatever shapes, shifts
/// and values the files held.
#[test]
fn every_load_check_of_the_networks_refuses_model_mismatch() {
use serde_json::json as j;
let m = Members::synthetic();
let conv = |o: &mut serde_json::Map<String, Value>, k: &str| -> serde_json::Map<String, Value> { o[k].as_object().unwrap().clone() };
let cases: Vec<(&str, &str, (Vec<u8>, bool))> = vec![
("a weight member missing", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o.remove("iact2.conv2.weight"); })),
("a layer norm's γ missing", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| { o.remove("TAM.TABs.1.ffn.norm1.weight"); })),
("a weight member of kind q16", "weights/hyp_y.json", edited_net("weights/hyp_y.json", |o| { let mut c = conv(o, "conv3.weight"); c.insert("kind".into(), j!("q16")); c.insert("v".into(), j!([1])); o.insert("conv3.weight".into(), Value::Object(c)); })),
("weights of other dims", "weights/hyp_uv.json", edited_net("weights/hyp_uv.json", |o| { o["conv1.weight"]["dims"] = j!([2, 95, 1, 1]); })),
("one weight too few", "weights/mcm1_y.json", edited_net("weights/mcm1_y.json", |o| { o["conv.0.weight"]["w"].as_array_mut().unwrap().pop(); })),
("a kernel that is not the topology's", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["TAM.ds_conv.weight"]["dims"] = j!([4, 4, 1, 1]); o["TAM.ds_conv.weight"]["w"] = j!(vec![1; 16]); })),
("a gate convolution whose groups are not 4", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["iact3.conv.weight"]["dims"] = j!([4, 2, 3, 3]); o["iact3.conv.weight"]["w"] = j!(vec![1; 72]); })),
("a context convolution whose groups are not 5", "weights/mcm2_y.json", edited_net("weights/mcm2_y.json", |o| { o["conv.1.weight"]["dims"] = j!([5, 5, 3, 3]); o["conv.1.weight"]["w"] = j!(vec![1; 225]); })),
("a bias one value short", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| { o["conv3_t.bias"]["v"].as_array_mut().unwrap().pop(); })),
("a shift of 128", "weights/mcm3_y.json", edited_net("weights/mcm3_y.json", |o| { o["fusion_pred_net.conv2.weight"]["shift"] = j!(128); })),
("a shift of -1", "weights/mcm0_y.json", edited_net("weights/mcm0_y.json", |o| { o["fusion_pred_net.conv1.weight"]["shift"] = j!(-1); })),
("a weight that is not an integer", "weights/hyp_y.json", edited_net("weights/hyp_y.json", |o| { o["conv2.weight"]["w"][0] = j!(0.5); })),
("a multiplier past i64", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["conv2_t.weight"]["mult"] = j!(9223372036854775808u64); })),
("three temperatures", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["TAM.TABs.0.attn.temperature"]["v"].as_array_mut().unwrap().pop(); })),
("ψ of other than 4C channels", "weights/hyp_uv.json", edited_net("weights/hyp_uv.json", |o| { o["conv4.weight"]["dims"] = j!([380, 2, 3, 3]); o["conv4.weight"]["w"] = j!(vec![1; 380 * 18]); o["conv4.bias"]["v"] = j!(vec![0; 380]); })),
("a luma synthesis of two planes", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["conv4_t.weight"]["dims"] = j!([4, 2, 3, 3]); o["conv4_t.weight"]["w"] = j!(vec![1; 72]); o["conv4_t.bias"]["v"] = j!([0, 0]); })),
("an odd feed-forward width", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| {
o["TAM.TABs.0.ffn.project_in.weight"]["dims"] = j!([3, 4, 1, 1]); o["TAM.TABs.0.ffn.project_in.weight"]["w"] = j!(vec![1; 12]);
o["TAM.TABs.0.ffn.dwconv.weight"]["dims"] = j!([3, 1, 3, 3]); o["TAM.TABs.0.ffn.dwconv.weight"]["w"] = j!(vec![1; 27]);
})),
("a member that is not an object of its kind", "weights/mcm1_y.json", edited_net("weights/mcm1_y.json", |o| { o["conv.1.bias"] = j!({ "kind": "q16", "v": "1" }); })),
// Past acc64: one weight of 2^40 makes a sum that no i64 holds under A32.
("a sum past i64", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["CAB.subscale.weight"]["w"][0] = j!(1i64 << 40); })),
];
let mut parsed_too = 0;
for (what, file, (b, typed)) in &cases {
let r = jpegai_ops::load_networks_bytes(nets_with(&m, file, b));
assert_eq!(r.as_ref().err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "{what}: the load from the bytes");
let parsed = Weights::from_json(&json(b));
assert_eq!(parsed.is_ok(), *typed, "{what}: this crate's parser reads exactly the typed files");
if !*typed {
parsed_too += 1;
} else {
let mut dw = synthetic_weights(&m);
let w = parsed.unwrap();
match *file {
"weights/hyp_y.json" => dw.luma.hyper_decoder = w,
"weights/hyp_uv.json" => dw.chroma.hyper_decoder = w,
"weights/synth_y.json" => dw.luma.synthesis = w,
"weights/synth_uv.json" => dw.chroma.synthesis = w,
f => dw.mcm_y[f.as_bytes()["weights/mcm".len()] as usize - b'0' as usize] = w,
}
let r = jpegai_ops::load_networks(&dw);
assert_eq!(r.as_ref().err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "{what}: the load from the parsed networks");
parsed_too += 1;
}
}
// Two members of one name: JSON this crate's parser reads as one member; the bytes are
// refused.
let file = "weights/hyp_y.json";
let text = String::from_utf8(m.get(file).to_vec()).unwrap();
let dup = text.replacen('{', "{\"conv3.bias\":{\"kind\":\"q16\",\"v\":[0,0]},", 1).into_bytes();
assert_eq!(jpegai_ops::load_networks_bytes(nets_with(&m, file, &dup)).err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "two members of one name");
// The committed bundle itself loads, from either form.
assert!(jpegai_ops::load_networks_bytes(m.nets()).is_ok() && jpegai_ops::load_networks(&synthetic_weights(&m)).is_ok());
assert_eq!(parsed_too, cases.len(), "every case is refused by the parser or the load from its result");
eprintln!("{} load checks refused model_mismatch from the bytes and from the parsed networks, and two members of one name", cases.len());
}
/// The committed synthetic bundle's networks as this crate parses them (with empty scale
/// networks and eICCI, which the networks' load does not read).
fn synthetic_weights(m: &Members) -> DecoderWeights {
let w = |n: &str| m.weights(n);
DecoderWeights {
luma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_y.json"))).unwrap(), hyper_decoder: w("weights/hyp_y.json"), synthesis: w("weights/synth_y.json") },
chroma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_uv.json"))).unwrap(), hyper_decoder: w("weights/hyp_uv.json"), synthesis: w("weights/synth_uv.json") },
mcm_y: [w("weights/mcm0_y.json"), w("weights/mcm1_y.json"), w("weights/mcm2_y.json"), w("weights/mcm3_y.json")],
icci: wai::jpegai_icci::IcciNet { luma: empty_path(), chroma: empty_path(), conv_last_v: empty_conv(), scale_v: 0.0 },
}
}
/// Gate 3, the ranges: a dequantised residual or a latent one past `A32`, at either end,
/// is refused `field_range` at its input, before any operator runs, by every graph that
/// reads it (on the committed networks); and, on the test codestream, a residual at the
/// range's end where the context stages' mean keeps the sum within it decodes, through
/// `rec_y` and `rec_uv`, to the direct decode's values, while one past it is refused
/// where the direct decode computed on (integer-operators §15.3 item 4).
#[test]
fn a_value_past_its_range_is_refused_field_range() {
use wai_ops::{Code, Site};
let m = Members::synthetic();
let syn = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
let past = [i32::MAX as i64 + 1, i32::MIN as i64 - 1];
let fill = |n: usize, at: usize, v: i64| -> Vec<i64> { (0..n).map(|i| if i == at { v } else { 0 }).collect() };
let mut refused = 0;
for v in past {
for (role, inputs, k) in [
("rec_y", vec![Tensor::new(160, 1, 1, vec![0; 160]), Tensor::new(160, 4, 4, fill(160 * 16, 7, v))], 1u32),
("rec_uv", vec![Tensor::new(96, 1, 1, vec![0; 96]), Tensor::new(96, 4, 4, fill(96 * 16, 95, v))], 1),
("syn_y", vec![Tensor::new(160, 2, 2, fill(640, 3, v))], 0),
("syn_uv", vec![Tensor::new(160, 2, 2, vec![0; 640]), Tensor::new(96, 2, 2, fill(384, 0, v))], 1),
] {
let r = wai_ops::run_with(syn.graph(role).unwrap(), inputs, &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host);
assert_eq!(r.err(), Some(wai_ops::Refusal::new(Code::FieldRange, Site::Input(k))), "{role}: input {k} holding {v}");
refused += 1;
}
}
eprintln!("{refused} inputs one past A32 refused at their input");
let Some(wh) = Whole::pinned("jpegai_nets_ops::a_value_past_its_range_is_refused_field_range") else {
return;
};
let i = inputs(&wh);
let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
for (k, plane) in [Plane::Luma, Plane::Chroma].into_iter().enumerate() {
// The mean each position gets: ŷ − r. A luma position is taken from the last
// sub-lattice, `(1, 0)`, which no context stage reads, so the other means stay.
let mu: Vec<i64> = i.yhat[k].iter().zip(&i.resi[k]).map(|(y, r)| y - r).collect();
let last = |p: usize| plane == Plane::Chroma || ((p % 256) / 16 % 2, p % 16 % 2) == (1, 0);
let lo_at = (0..mu.len()).find(|&p| last(p) && mu[p] >= 0).unwrap();
let hi_at = (0..mu.len()).find(|&p| last(p) && mu[p] <= 0).unwrap();
for (at, edge, one_past) in [(lo_at, i32::MIN as i64, i32::MIN as i64 - 1), (hi_at, i32::MAX as i64, i32::MAX as i64 + 1)] {
let mut r = i.resi[k].clone();
r[at] = edge;
let got = jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), r.clone()).unwrap_or_else(|e| panic!("{plane:?}: a residual of {edge} where the mean is {}: {e:?}", mu[at]));
let want: Vec<i64> = r.iter().zip(&mu).map(|(a, b)| a + b).collect();
assert!(got == want, "{plane:?}: a residual at {edge} decodes as the direct decode would");
r[at] = one_past;
assert_eq!(jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), r).err(), Some(wai_ops::Refusal::new(Code::FieldRange, Site::Input(1))), "{plane:?}: a residual of {one_past}");
}
}
eprintln!("on the test codestream: a residual at either end of A32 decodes as directly, one past it is refused field_range");
}

40
wai/wai-rs/tests/jpegai_ops.rs
View file

@ -24,7 +24,7 @@ use wai::jpegai_synthesis::IntNet;
use wai_ops::Profile;
const CAP: &str = "wai.image.jpegai";
const SYNTHETIC: &str = "60020d0657c4561db336a42f2d359b254a3f12348e8122aeb7bbb91d63b180ad";
const SYNTHETIC: &str = "1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9";
const V1: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37";
const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)";
@ -91,7 +91,7 @@ impl Bundle {
Some((Bundle { norm: read(Path::new(&norm)), hs_y: read(&d.join("hs_y.json")), hs_uv: read(&d.join("hs_uv.json")) }, d))
}
fn parsed(&self) -> (NormTables, IntNet, IntNet) {
(norm_tables(&json(&self.norm)), IntNet::from_json(&json(&self.hs_y)), IntNet::from_json(&json(&self.hs_uv)))
(norm_tables(&json(&self.norm)), IntNet::from_json(&json(&self.hs_y)).unwrap(), IntNet::from_json(&json(&self.hs_uv)).unwrap())
}
}
@ -99,9 +99,10 @@ impl Bundle {
/// converted (`jpegai_ops::net`), is the bundle `wai-ops` parses from those bytes and
/// lowers to the same graphs; the scale path this crate loads (`jpegai_ops::load`) holds
/// those graphs and meets `acc32`; and the committed certificate `digest` is those graphs',
/// with the verdict Appendix C predicts (`acc32`, in the accelerator subset, no flags).
/// each of whose nodes meets `acc32` with no accelerator flag (the scale path's
/// `role_profiles`, §11).
fn lowers_to_its_certificate(name: &str, b: &Bundle, digest: &str) {
let lowered = wai_ops::lower::jpegai(&b.norm, &b.hs_y, &b.hs_uv).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
assert_eq!(roles, ["hs_y", "hs_uv"], "{name}");
let (n, y, uv) = b.parsed();
@ -121,23 +122,32 @@ fn lowers_to_its_certificate(name: &str, b: &Bundle, digest: &str) {
let x = &cert["graphs"][i]["inputs"][0];
assert_eq!((x["policy"].as_str(), x["lo"].as_str(), x["hi"].as_str()), (Some("check"), Some("-31"), Some("32")), "{name}: z's range");
}
let v = &cert["verdict"];
assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool(), v["npu"].as_bool()), (Some(true), Some(true), Some(true)), "{name}: the certificate's verdict");
assert_eq!(v["npu_flags"].as_array().map(Vec::len), Some(0), "{name}: no accelerator flags");
// The set's verdict covers its other networks too (§13.6.2), which meet acc64 only;
// the scale path's two graphs meet acc32, in the accelerator subset, with no flags
// (§11's `role_profiles`).
assert_eq!(cert["verdict"]["acc64"].as_bool(), Some(true), "{name}: the certificate's verdict");
for i in 0..2 {
for n in cert["graphs"][i]["nodes"].as_array().unwrap() {
assert!(n["fits"].as_array().unwrap().iter().any(|p| p == "acc32"), "{name}: node {} of the scale path fits acc32", n["index"]);
assert_eq!(n["npu_flags"].as_array().map(Vec::len), Some(0), "{name}: no accelerator flags");
}
}
}
/// The committed synthetic bundle: its lowered graphs are the committed ones, and its
/// set digest is the listing of its three members.
/// The committed synthetic bundle: its lowered scale graphs are the committed ones, and
/// its set digest is the listing of its eleven members.
#[test]
fn the_committed_bundle_lowers_to_its_committed_graphs_and_verdict() {
let b = Bundle::synthetic();
let dir = wai().join(format!("ops-conformance/certificates/{CAP}/{SYNTHETIC}"));
let lowered = wai_ops::lower::jpegai(&b.norm, &b.hs_y, &b.hs_uv).unwrap();
let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap();
for role in ["hs_y", "hs_uv"] {
assert!(lowered.get(role).unwrap().to_bytes() == read(&dir.join(format!("{role}.whg1"))), "the lowering of the bytes differs from the committed {role}.whg1");
}
let members: BTreeMap<String, &[u8]> =
let nets: Vec<(String, Vec<u8>)> = wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().map(|n| (n.to_string(), read(&wai().join("ops-conformance/sets/jpegai").join(n)))).collect();
let mut members: BTreeMap<String, &[u8]> =
[("norm.json", &b.norm), ("weights/hs_y.json", &b.hs_y), ("weights/hs_uv.json", &b.hs_uv)].into_iter().map(|(k, v)| (k.to_string(), v.as_slice())).collect();
members.extend(nets.iter().map(|(k, v)| (k.clone(), v.as_slice())));
let digest = wai::int_set_cache::SetKey::listing_digest(&members).unwrap();
assert_eq!(digest.iter().map(|x| format!("{x:02x}")).collect::<String>(), SYNTHETIC, "the set's digest");
lowers_to_its_certificate("synthetic", &b, SYNTHETIC);
@ -174,13 +184,13 @@ impl Whole {
let [cs, dir, nf] = require_vectors::env_paths(test, ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
let d = PathBuf::from(dir);
let j = |n: &str| json(&read(&d.join(n)));
let w = |n: &str| Weights::from_json(&j(n));
let net = |n: &str| IntNet::from_json(&j(n));
let w = |n: &str| Weights::from_json(&j(n)).unwrap();
let net = |n: &str| IntNet::from_json(&j(n)).unwrap();
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")),
icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")).unwrap(),
};
Some(Whole { data: read(Path::new(&cs)), weights, norm: norm_tables(&json(&read(Path::new(&nf)))) })
}
@ -296,7 +306,7 @@ fn every_load_check_of_the_bundle_refuses_model_mismatch() {
// Past acc32 and within acc64: block 3's sums exceed i32 with one weight of 2^30 over
// inputs in [0, 127].
let big = edited(|_, y| { y["pointwise.weight"]["v"][0] = Value::from(1i64 << 30); });
let g = wai_ops::lower::jpegai(&big.0.norm, &big.0.hs_y, &big.0.hs_uv).expect("the bundle lowers");
let g = wai_ops::lower::jpegai_scale(&big.0.norm, &big.0.hs_y, &big.0.hs_uv).expect("the bundle lowers");
assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc64).is_ok(), "it meets acc64");
assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc32).is_err(), "and not acc32");
cases.push(("hs_y: a sum past i32 (acc64, not acc32)", big));

4
wai/wai-rs/tests/jpegai_reconstruct_calibration.rs
View file

@ -10,7 +10,7 @@
//!
//! env: WAI_JPEGAI_RECON_FIXTURE = recon_fixture.json {resi_dq_y, psi_y, yhat_y, c, h2, w2}
//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
#![cfg(feature = "neural_int")]
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
use wai::jpegai_pipeline::reconstruct_yhat;
use wai::jpegai_synthesis::Weights;
@ -41,7 +41,7 @@ fn reconstruct_yhat_on_real_reference_data() {
let load = |i: usize| -> Weights {
let raw = std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap();
Weights::from_json(&serde_json::from_slice(&raw).unwrap())
Weights::from_json(&serde_json::from_slice(&raw).unwrap()).unwrap()
};
let (w0, w1, w2w, w3) = (load(0), load(1), load(2), load(3));

6
wai/wai-rs/tests/jpegai_reconstruction_binding.rs
View file

@ -58,13 +58,13 @@ fn real_decode_to_reconstruction_binding() {
// --- load the model bundle (same shape as jpegai_e2e_full) ---
let dpath = std::path::Path::new(&dir);
let w = |n: &str| Weights::from_json(&jload(&dpath.join(n)));
let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n)));
let w = |n: &str| Weights::from_json(&jload(&dpath.join(n))).unwrap();
let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n))).unwrap();
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
icci: IcciNet::from_json(&jload(&dpath.join("icci.json"))),
icci: IcciNet::from_json(&jload(&dpath.join("icci.json"))).unwrap(),
};
let nv = jload(std::path::Path::new(&nf));
let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();

4
wai/wai-rs/tests/jpegai_synthesis_calibration.rs
View file

@ -6,7 +6,7 @@
//!
//! env: WAI_JPEGAI_SYNTH16_Y = path to synthesis16.onnx (model_y),
//! WAI_JPEGAI_Y_WEIGHTS = path to the exported 16-bit weights JSON.
#![cfg(feature = "neural")]
#![cfg(all(feature = "neural", feature = "legacy_int"))]
use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
@ -46,7 +46,7 @@ fn synthesis_int_vs_float_psnr() {
let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
let yi = decode_y(&li, &wt);
assert_eq!(yi.len(), yf.len(), "output size mismatch");

4
wai/wai-rs/tests/jpegai_uv_calibration.rs
View file

@ -5,7 +5,7 @@
//!
//! env: WAI_JPEGAI_SYNTH16_UV = path to model_uv/synthesis16.onnx,
//! WAI_JPEGAI_UV_WEIGHTS = path to the exported 16-bit uv weights JSON.
#![cfg(feature = "neural")]
#![cfg(all(feature = "neural", feature = "legacy_int"))]
use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
@ -45,7 +45,7 @@ fn synthesis_uv_int_vs_float_psnr() {
let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
let wt = Weights::from_json(&v);
let wt = Weights::from_json(&v).unwrap();
let uvi = decode_uv(&li, &wt);
assert_eq!(uvi.len(), uvf.len(), "output size mismatch (2×256×256 expected)");

19
wai/wai-verify/src/bin/wai-verify/corpus.rs
View file

@ -265,6 +265,25 @@ fn check_external(s: &json::Value, root: &Path) -> Result<(), String> {
if verdict.get(profile.as_str()).and_then(|x| x.as_bool()) != Some(true) {
return Err(format!("{path}: its verdict does not meet {}", profile.as_str()));
}
// A role held to a stronger profile (§11, `role_profiles`): that graph's inputs and
// scalars meet the profile's tensor width, and each of its nodes fits it (§8.5).
if let Some(roles) = s.get("role_profiles").and_then(|x| x.as_obj()) {
let graphs = v.get("graphs").and_then(|x| x.as_arr()).ok_or("graphs")?;
for (role, p) in roles {
let p = p.as_str().and_then(Profile::parse).ok_or("role_profiles")?;
let g = graphs.iter().find(|g| g.get("name").and_then(|x| x.as_str()) == Some(role)).ok_or_else(|| format!("{path} has no graph {role}"))?;
let width = if p == Profile::Acc32 { 32 } else { 64 };
let fits = |r: &json::Value| -> bool {
let n = |k: &str| r.get(k).and_then(|x| x.as_dec_i128());
matches!((n("lo"), n("hi")), (Some(lo), Some(hi)) if lo >= -(1i128 << (width - 1)) && hi < 1i128 << (width - 1))
};
let ranges = ["inputs", "scalars"].iter().all(|k| g.get(k).and_then(|x| x.as_arr()).is_some_and(|a| a.iter().all(fits)));
let nodes = g.get("nodes").and_then(|x| x.as_arr()).is_some_and(|a| a.iter().all(|n| n.get("fits").and_then(|f| f.as_arr()).is_some_and(|f| f.iter().any(|x| x.as_str() == Some(p.as_str())))));
if !(ranges && nodes) {
return Err(format!("{path}: its graph {role} does not meet {}", p.as_str()));
}
}
}
Ok(())
}

27
wai/wai-verify/src/bin/wai-verify/main.rs
View file

@ -4,6 +4,7 @@
//! ```text
//! wai-verify bounds <graph.whg1>... | <set-path> [--capability <name>] [--lowering <id>]
//! [--member <name>=<path>]... [--c <n>] [--profile acc32|acc64]
//! [--role-profile <role>=acc32|acc64]...
//! [--out <cert.json>] [--check <committed.json>]
//! [--advise acc32|acc64] [--witness <dir>]
//! wai-verify lower <set-path> --capability <name> [--member <name>=<path>]... [--c <n>]
@ -19,12 +20,16 @@
//! every lowering of §13; `--c` is `world/1`'s latent channel count. `WHG1` graphs
//! are certified on their own, in the order given, with the lowering `"whg1"`.
//!
//! `--role-profile` holds one graph of the set, by its role, to a profile of its own
//! (§11, `role_profiles`): the graph's inputs, scalars and nodes must meet it.
//!
//! The certificate (§8.5) goes to `--out`, or to stdout. The verdict, the first
//! failure, the witnesses' paths and the advice go to stderr.
//!
//! Exit status:
//! - 0: the certificate was written and its verdict meets the profile;
//! - 1: the verdict fails the profile (the certificate is still written);
//! - 1: the verdict fails the profile, or a graph `--role-profile` names fails its
//! profile (the certificate is still written);
//! - 2: malformed input: a graph refused at load steps 1–7, an empty `check`
//! intersection or a `lut16` range past `[−32767, 32767]`, or an unreadable file;
//! - 3: `--check` found a difference, printed as the first differing JSON path;
@ -46,6 +51,7 @@ use wai_ops::json::{self, C, Value};
const USAGE: &str = "usage:
wai-verify bounds <graph.whg1>... | <set-path> [--capability <name>] [--lowering <id>]
[--member <name>=<path>]... [--c <n>] [--profile acc32|acc64]
[--role-profile <role>=acc32|acc64]...
[--out <cert.json>] [--check <committed.json>]
[--advise acc32|acc64] [--witness <dir>]
wai-verify lower <set-path> --capability <name> [--member <name>=<path>]... [--c <n>] --out <dir>
@ -92,6 +98,7 @@ struct Args {
members: Vec<(String, String)>,
c: Option<u32>,
profile: Option<Profile>,
role_profiles: Vec<(String, Profile)>,
out: Option<String>,
check: Option<String>,
advise: Option<Profile>,
@ -114,6 +121,11 @@ fn parse_args(a: &[String]) -> Result<Args, String> {
}
"--c" => r.c = Some(val(k)?.parse().map_err(|_| "--c takes a channel count")?),
"--profile" => r.profile = Some(Profile::parse(&val(k)?).ok_or("--profile is acc32 or acc64")?),
"--role-profile" => {
let v = val(k)?;
let (n, p) = v.split_once('=').ok_or("--role-profile takes <role>=acc32|acc64")?;
r.role_profiles.push((n.to_string(), Profile::parse(p).ok_or("--role-profile takes <role>=acc32|acc64")?));
}
"--advise" => r.advise = Some(Profile::parse(&val(k)?).ok_or("--advise is acc32 or acc64")?),
"--out" => r.out = Some(val(k)?),
"--check" => r.check = Some(val(k)?),
@ -328,6 +340,19 @@ fn cmd_bounds(args: &Args) -> ExitCode {
}
eprintln!("npu: {} (reasons {:?}, flags {:?})", v.npu(), v.npu_reasons, v.npu_flags);
let mut code = if v.meets(profile) { 0 } else { 1 };
for (role, p) in &args.role_profiles {
match cert.graphs.iter().find(|cg| cg.name == *role) {
Some(cg) if cg.bounds.meets(&cg.graph, *p) => eprintln!("{role}: meets {}", p.as_str()),
Some(_) => {
eprintln!("{role}: fails {}", p.as_str());
code = 1;
}
None => {
eprintln!("wai-verify bounds: the set has no graph {role}");
return ExitCode::from(5);
}
}
}
if let Some(dir) = &args.witness
&& !v.meets(profile)
{

7
wai/wai-web/src/neural_int.rs
View file

@ -263,14 +263,15 @@ pub fn decode_jpegai_full(
let parse = |b: &[u8]| -> Result<serde_json::Value, JsError> {
serde_json::from_slice(b).map_err(|e| JsError::new(&e.to_string()))
};
let w = |b: &[u8]| -> Result<Weights, JsError> { Ok(Weights::from_json(&parse(b)?)) };
let net = |b: &[u8]| -> Result<IntNet, JsError> { Ok(IntNet::from_json(&parse(b)?)) };
let bad = |e: wai::jpegai_synthesis::BadMember| JsError::new(&format!("wai.image.jpegai: {e}"));
let w = |b: &[u8]| -> Result<Weights, JsError> { Weights::from_json(&parse(b)?).map_err(bad) };
let net = |b: &[u8]| -> Result<IntNet, JsError> { IntNet::from_json(&parse(b)?).map_err(bad) };
let weights = DecoderWeights {
luma: PlaneWeights { hyper_scale: net(hs_y)?, hyper_decoder: w(hyp_y)?, synthesis: w(synth_y)? },
chroma: PlaneWeights { hyper_scale: net(hs_uv)?, hyper_decoder: w(hyp_uv)?, synthesis: w(synth_uv)? },
mcm_y: [w(mcm0)?, w(mcm1)?, w(mcm2)?, w(mcm3)?],
icci: IcciNet::from_json(&parse(icci)?),
icci: IcciNet::from_json(&parse(icci)?).map_err(bad)?,
};
let nv = parse(norm)?;