Skip to content

feat(compression): update tooling to use DECODE operators - #3647

Draft
rkuester wants to merge 16 commits into
tensorflow:mainfrom
rkuester:feat-decode/queue
Draft

feat(compression): update tooling to use DECODE operators#3647
rkuester wants to merge 16 commits into
tensorflow:mainfrom
rkuester:feat-decode/queue

Conversation

@rkuester

@rkuester rkuester commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

This is a draft PR for running CI, review, and seeing the commits in
context. The commits along this branch will be individually submitted
for merge.

This replaces #3574, whose head branch was renamed to
feat-decode/queue for a stacked-commits workflow; renaming a fork
branch closes its cross-repo PR, so the draft is reopened here with
the same content under the new head.

See the linked issue for a description of the change.

BUG=implements #3256

@rkuester
rkuester force-pushed the feat-decode/queue branch 2 times, most recently from 266a467 to 34de4fa Compare August 12, 2026 22:26
@rkuester
rkuester force-pushed the feat-decode/queue branch 5 times, most recently from d3133d9 to 81d096f Compare August 20, 2026 19:19
rkuester added 16 commits August 25, 2026 16:29
@rkuester
docs(python): explain env vars in test runner
a70c6ae 
Comment the two environment variables set before the tests run.

BUG=part of tensorflow#3256
@rkuester
fix(compression): fix .tflite built from scratch
901cc1b 
A model built from scratch with the model editor declares schema
version 0 and lacks the TFL3 file identifier. Loaders that verify
these fields, such as the TfLite interpreter, reject such a model.
Set version 3 on a model that declares no version, and finish every
built flatbuffer with the file identifier. Leave the version of a
model read from a file unchanged.
@rkuester
fix(compression): set operator index on read and build
2da8f5b 
Reading or building a model sets the index of every tensor and
subgraph, but an operator's index is set only when the operator is
added imperatively to a subgraph. Operators of a model read from a
flatbuffer report no index. Set the operator index in the read and
build paths too, and document when the index is valid, matching the
tensor and subgraph indexes.
@rkuester
feat(compression): add spec proposal tool
d343607 
Writing a compression spec by hand means searching the model for the
tensors worth compressing and working out a bitwidth for each. Add a
tool that walks a model and proposes a spec entry for every constant
tensor that LUT compression can encode, with the minimum
index_bitwidth able to enumerate the tensor's unique values. Comment
each entry with the tensor's identity and its size arithmetic, so a
reviewer can read the proposal and delete the entries for tensors
that should stay uncompressed:

  # "arith.constant" INT32 [4, 2], input 1 of PAD (operators 170, 2446)
  # 8 elements, 2 unique values, 32 -> 25 bytes (7 saved, 22%)
  - subgraph: 0
    tensor: 36
    compression:
      - lut:
          index_bitwidth: 1

Propose by default only the tensors that compression would shrink,
counting the packed indices, the value tables, and the decode header.
Add a flag that widens the proposal to every encodable constant. In a
footer, list the constants left out and why. Annotate tensors that
share a buffer, where the converter deduplicated identical constants,
so a reviewer keeps or deletes all aliases together. Estimate the
whole model file's size change in a header line.
@rkuester
refactor(compression): name operators in the model editor
42e675e 
Give the operator wrapper a property returning its kind as text.
Custom operators go by their custom code, builtins by the name of
their enumerator, and an unrecognized code by its number.

The spec proposal tool derived this from the schema enum itself.
Naming an operator is a property of an operator, so it belongs with
the wrapper that models one, where every caller can reach it
instead of repeating the lookup.
@rkuester
feat(compression): add explicit compression mode
cb307f0 
Let a compression spec state whether a tensor gets one value
table or one table per channel along a given axis, rather than
leaving the choice to inference from the tensor's quantization.
Honor an explicit choice in the LUT compressor and record it in a
new axis field in the header byte that holds the bitwidth: values
0--14 name the channel axis of the output tensor's shape, and 15
means one table. Reject axes the kernels do not support. Current
kernels mask off the axis field and infer the table layout from
quantization, so a model whose explicit choice agrees with its
quantization runs on them unchanged.

A spec without an explicit choice compresses by inference as
before and leaves the axis field zero. This path is
transitional, kept so the change merges in small reviewable
commits. The spec-file parser does not yet accept the choice; a
later commit adds parsing, requires the choice, and removes the
inference.
@rkuester
feat(compression): build specs with explicit mode
ce36450 
Extend the spec builder's fluent API so a lookup-table entry can
state the compression mode, per tensor or per channel along a
given axis:

    SpecBuilder().add_tensor(subgraph=0, tensor=2) \
        .with_lut(index_bitwidth=4, mode=PerChannel(axis=0)) \
        .build()

Omitting the mode leaves the choice to inference from the
tensor's quantization, the transitional default, until the
spec-file parser accepts and requires the choice in an upcoming
commit.
@rkuester
feat(compression): propose the compression mode
59424fd 
Write the per-tensor or per-channel choice into each entry of a
proposed compression spec, drawn from the tensor's quantization:
per_channel with the quantized axis for a per-channel-quantized
tensor, per_tensor otherwise. The proposal already computed the
choice to size the value tables; writing it out shows it to the
human reviewing the proposal. The spec-file parser ignores the
keys today; a later commit makes them required, and proposals
written from now on parse unchanged across that cut.
@rkuester
feat(compression): require the mode in spec files
828083a 
Require each lut entry of a compression spec file to give exactly
one of per_tensor or per_channel. per_channel takes a mapping
with a non-negative integer axis; per_tensor takes no value.
Parsing fails, naming the two choices, for an entry giving
neither or both, and for a malformed value under either key.

The compressor still tolerates a spec object built without a
mode; the next commit removes that tolerance along with the
inference from quantization it falls back on.
@rkuester
feat(compression): drop mode inference
10c8f0c 
Remove the transitional path that inferred a tensor's compression
mode from its quantization when a spec object gave none. A spec
object without a mode is now an error, matching the spec-file
parser, which already requires the choice, so the header's axis
field is always written explicitly. State the mode in every test
fixture's spec; each states what its tensor's quantization
implied, so the compressed bytes are unchanged except the axis
field.

Move the quantization-based axis inference into the proposal
tool, its only remaining caller, where it drafts the proposed
mode for human review.
@rkuester
feat(compression): choose the mode from the values
8fdd479 
Choose a proposed tensor's per-tensor or per-channel mode by
testing both layouts against its values and proposing the one
that encodes the tensor smallest. Stop reading the mode from
the tensor's quantization.

Quantization says nothing about how weights were binned, and
an unquantized tensor has none to read. Weights binned per
channel hold few distinct values within a channel and many
across the whole tensor, so the old rule proposed one value
table, found more distinct values than the index can
enumerate, and dropped the tensor as unencodable. The weights
that compress best were the ones left out.

When no layout fits, name the closest one tried, so a
rejection shows whether a channel axis would have helped.

Skip a channel axis whose slices hold a single element. Every
element would go in a value table and the indices would be
pure overhead.
@rkuester
test(compression): tighten per-channel tests
fded12b 
Give the per-channel integration model distinct values in each
channel, so a decode that used the wrong table or stride no
longer produces the right output by accident. Assert the packed
value tables and the channel stride byte by content in the
compressor's unit tests. Add an end-to-end test of a tensor
quantized along its last axis, which nothing ran through the
compressor and interpreter together before.

Add an end-to-end test of a per-channel choice on an unquantized
tensor, which the current kernel decodes wrongly because it
derives the channel layout from quantization and ignores the
header's axis field. Mark the test as an expected failure until
the kernel reads the axis field.
@rkuester
docs(compression): update the spec-file example
3b07907 
Add the compression mode, now required, to each entry of the
example spec file, and describe the two forms: per_channel with
the axis that gives the channel count, and a bare per_tensor.
@rkuester
feat(compression): stop proposing required constants
26327a3 
Leave a tensor whose consumer must read it during Prepare out of a
proposed spec, naming the kernel and the input position that
require it. Compression replaces a constant with the output of a
DECODE operator, and that output holds no values until DECODE runs
in Invoke, so the model fails to prepare.

List the inputs every implementation requires, keyed by operator
and input position, in one table. The compressor shares the table
in the commit that follows, which refuses such a tensor named in a
spec written by hand.
@rkuester
feat(compression): refuse to compress required constants
846d5f5 
Refuse to compress a tensor whose consumer must read it during
Prepare, naming the kernel and the input position that require it.
Compression replaces a constant with the output of a DECODE
operator, and that output holds no values until DECODE runs in
Invoke, so the model fails to prepare.

A proposed spec already leaves such a tensor out. A spec written or
edited by hand can still name one, and until now the compressor
took it.
@rkuester
feat(compression): add a minimum savings floor
df01570 
Take a --min_savings floor in bytes and leave out any tensor whose
compression saves less, naming the shortfall in the footer beside
the other reasons a tensor is left out. Replace the boolean
--require_savings, whose off position becomes --min_savings=0 and
whose on position becomes any positive floor.

The old rule listed a tensor whenever compression shrank it by a
single byte. That undercounts the cost. Compressing a tensor also
adds a DECODE operator, an ancillary tensor, and a decoded output
tensor, and the byte estimate counts none of them. Measured on a
production model, that structure runs 120 to 300 bytes per entry,
varying with the length of the tensor names, so an entry saving
less than that grows the model it was meant to shrink.

Default the floor to 512 bytes, above the measured range, so a
proposed entry pays for itself. Nothing is hidden by the default.
A spec is a draft for review, and a tensor held back by the floor
still appears in the footer with the bytes it would have saved, so
a reader who disagrees can lower the floor and get it back.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@rkuester