feat(compression): update tooling to use DECODE operators#3574
Draft
rkuester wants to merge 29 commits into
Draft
feat(compression): update tooling to use DECODE operators#3574rkuester wants to merge 29 commits into
rkuester wants to merge 29 commits into
Conversation
This was referenced May 26, 2026
rkuester
force-pushed
the
feat-decode
branch
6 times, most recently
from
June 1, 2026 22:01
a0e3ff8 to
75ca9dc
Compare
rkuester
force-pushed
the
feat-decode
branch
2 times, most recently
from
June 30, 2026 00:47
7f5e8cb to
41867cb
Compare
rkuester
force-pushed
the
feat-decode
branch
2 times, most recently
from
July 6, 2026 20:41
e09ab22 to
cb8e6f1
Compare
Insert DECODE operators before consumers of compressed tensors. Each consumer gets its own DECODE operator to support alternate decompression memory, which resets allocations between DECODE invocations. After insertion, compressed tensors are rewritten to hold encoded data as UINT8 with shape matching byte count. BUG=part of tensorflow#3256
A compressed tensor can be listed in a subgraph's output list, where it is read not by an operator, but by the operator calling the subgraph (IF, WHILE), which copies subgraph outputs when the subgraph returns, or by the client, which reads model outputs after invocation. DECODE insertion previously checked the output list only for tensors with no consumers, and refused those as unsupported. A tensor both consumed and listed as an output slipped through. The pass rewired its consumers to decoded values, then rewrote the tensor to hold encoded bytes, which the output list delivered as if decoded. Treat the output list as one more consumer, one which reads its tensors only after the last operator runs. Append a DECODE after the last operator for each compressed tensor in the output list, and rewire the list entry to the decoded value. BUG=part of tensorflow#3256
A consumer reading several compressed tensors needs all their decoded values at once, but under alternate decompression memory, values produced by different DECODE operators cannot coexist. Each DECODE resets the allocation offset during Prepare, placing every DECODE's outputs at the same address, so each DECODE overwrites the outputs of the one before it. Decoding a consumer's tensors with separate DECODE operators corrupts all but the last value. Decode all compressed tensors read by one consumer with a single DECODE operator carrying one encoded/ancillary input pair and one output per tensor. Outputs of a single DECODE coexist, since the allocation reset happens between operators, not between the outputs of one. The subgraph output list, treated as one more consumer, gets the same treatment. One DECODE, appended after the last operator, decodes every compressed tensor in the list. BUG=part of tensorflow#3256
DECODE insertion sorted consumers and located insertion points with list.index, a linear scan of the operator list per lookup. Build a map of operator positions once per subgraph and consult it instead. The positions recorded before any insertion remain correct throughout, because consumers are handled in reverse position order, so each insertion falls after every consumer still to be processed. BUG=part of tensorflow#3256
Add a test suite that exercises DECODE outputs crossing subgraph boundaries on the TFLM interpreter, rather than only checking the rewritten flatbuffer structure. The tests build multi-subgraph WHILE models with model_editor, compress constants with the LUT compressor, insert DECODE operators with decode_insert, and verify inference results, in both arena and alternate decompression memory modes. The case of a DECODE output feeding a WHILE input, with a second DECODE in the cond subgraph and alternate decompression memory in use, requires the WhileEval fix from tensorflow#3633 (issue tensorflow#3632). WHILE formerly re-read its inputs after invoking the cond subgraph, picking up the value the cond subgraph's DECODE wrote over shared alternate memory.
Add Tensor.copy(), which duplicates a tensor's backing TensorT and shares the original's Buffer object. Duplicating the TensorT preserves fields model_editor does not otherwise manage, such as is_variable and shape_signature. An optional name argument gives the copy its own name. Clients that need a data-less copy, such as tooling that creates stand-in tensors, can assign None to the copy's buffer. Add Tensor.equal(), a field-wise equality over the backing TensorT, quantization, and buffer. Fields unknown to model_editor participate via recursive comparison, so clients can compare tensors without enumerating fields. Buffers compare by identity, mirroring how the model expresses buffer sharing.
Add dedupe_buffers(), which repoints tensors whose buffers hold byte-identical contents at one canonical Buffer object, mirroring the TfLite converter's deduplication of identical constants. Tensors marked is_variable are left alone, since mutable data must not alias. The walk covers every tensor the compiler collects, including tensors inline on operators that never appear in a subgraph's tensor list. Merged-away buffers linger in model.buffers until pruned.
Add prune_buffers(), which rebuilds model.buffers with only the conventional empty buffer 0 and the buffers some tensor references, renumbering indices in the process. Models built from scratch keep an empty buffer list and compile only referenced buffers, so pruning matters for models from read(), whose buffer list the compiler preserves wholesale, including entries orphaned by editing.
Create the output tensor of a DECODE operator by copying the original tensor, clearing its data, and renaming it, rather than by building a new tensor from the original's shape, dtype, and quantization. Copying preserves TensorT fields the insertion code does not otherwise handle, such as is_variable and shape_signature, so the decoded stand-in is indistinguishable from the original tensor it replaces. Verify the output against a copy of the original snapshotted before insertion rewrites it, compared with field-wise tensor equality so every field participates without the test enumerating them.
Distinct tensors can share one buffer, in the same or different subgraphs, where the converter deduplicated identical constants. Give each rewritten encoded tensor and each ancillary tensor a fresh buffer, then merge byte-identical buffers and prune unreferenced ones after insertion. Sharing survives compression wherever aliases compress to identical results, extends to any ancillary data that coincides, and dissolves where results diverge (possible for tensors sharing bytes but quantized with different structures), rather than one alias corrupting another through a shared buffer rewritten in place. Skip, with a warning, compressed tensors that share a buffer with an uncompressed tensor. The uncompressed data must remain in the model for the other tensors, so compressing such an alias cannot reduce model size.
Describe the placement of DECODE operators in terms of the operator's contract. Outputs have a lifetime limited to the very next operator in the subgraph, and DECODE trades increased latency for decreased memory usage. Remove the explanation of interpreter alternate-memory behavior that previously justified the per-consumer placement, along with a confusing aside about clients reading model outputs. Describe the output tensor as a copy of the original, matching the implementation.
Add Subgraph.consumers_of(), which returns the operators reading a given tensor, in subgraph order. Replace decode_insert's private helper _find_tensor_consumers with it. The helper's unit test called a private method of decode_insert; consumer lookup is now a public API, tested in model_editor_test.
Specify the subgraph inputs and outputs in the three test model builders. Two builders previously declared neither, and the third declared only its outputs, so the subgraph compiler emitted empty vectors for whatever was missing, making the models structurally unlike anything the converter produces.
Remove two weights_tensor assignments never read by their tests.
When one compressed tensor feeds multiple DECODE operators, assert that the operators read the same encoded tensor object, alongside the existing assertion that they share the ancillary tensor. Cover both situations in which one tensor feeds multiple DECODE operators, a tensor with two consumers and a tensor both consumed and listed as a subgraph output.
Compare the DCM's decode type byte against the DecodeType.LUT constant instead of a magic zero. Convert the DCM slice from a numpy array to bytes first, so that indexing yields a plain integer whose comparison defers to the constant's own equality.
Exercise the one-DECODE-per-consumer batching with a CONCATENATION of two compressed tensors instead of a FULLY_CONNECTED with an extra weights input, which is not a valid FC signature. CONCATENATION takes any number of inputs, so the model resembles something a converter could produce.
Replace the fixed dummy ancillary data helper with one that builds a whole CompressionResult, parameterized by element count, index bitwidth, and value table, so each test passes values consistent with the tensor it compresses and the encoded data is sized accordingly. State in the module docstring that the test models and payloads are structural fixtures, not valid runnable models or decodable data, so a reader does not mistake them for real examples.
Two tensors can share one buffer when the converter deduplicates identical constants. Add tests for the two situations in which insertion cannot preserve that sharing. In the first, both tensors are compressed but their compression results differ. Insertion dissolves the sharing, and the test verifies that each tensor receives its own encoded and ancillary buffers holding its own results. In the second, only one of the tensors is compressed. The uncompressed tensor keeps the original data in the model, so compressing its alias would grow the model rather than shrink it. Insertion declines to compress, and the test verifies that no DECODE operator is inserted, that the tensor is untouched, and that a warning explains why.
All other insertion tests build their models from scratch, and a from-scratch model keeps an empty buffer list, which makes buffer pruning a no-op. Add a test that round-trips a model through build() and read() before insertion, then verifies the packed result carries the DECODE operator with its encoded and ancillary data, and that the buffer orphaned when compression rewrote the weights tensor is pruned rather than left in the model.
Replace monolithic compression logic with a dispatch table that routes compression requests to plugin modules based on the spec's compression method type. After compressing tensors, insert DECODE operators into the model graph. Warn when compression expands data, helping users identify tensors that don't benefit from compression. BUG=part of tensorflow#3256
Add tests that compress models with LUT compression, run them through the TFLM Python interpreter, and verify outputs match uncompressed originals. Cover per-tensor and per-channel quantization, various index bitwidths, unquantized weights, and alternate decompression memory. BUG=part of tensorflow#3256
Add a manual test for verifying compression on proprietary models that can't be checked into the repository. See the module docstring for usage instructions. BUG=part of tensorflow#3256
Explicit inheritance from Protocol enables static type checking at definition time and makes the interface self-documenting. BUG=part of tensorflow#3256
An upcoming change registers the DECODE operator unconditionally in the Python ops resolver, after which compress() emits DECODE-based models that load successfully. That breaks this test's original approach, which ran a model through compress() and expected the load to fail. Rewrite it to instead inject a raw COMPRESSION_METADATA entry into the flatbuffer via model_editor, directly exercising the HasCompressionMetadata() detection path for legacy-compressed models. Decoupling the test from compress() output lets it verify the legacy-rejection behavior independently of whether the DECODE operator is registered, so it passes both before and after that upcoming change. BUG=part of tensorflow#3256
The DECODE kernel and its dependencies are already compiled unconditionally -- none are guarded by USE_TFLM_COMPRESSION. Remove the #ifdef around AddDecode() in PythonOpsResolver so DECODE-based compressed models work in a default Python build. Remove the with_compression_enabled gating from compression and proprietary integration tests, since they use DECODE-based models that no longer require the flag. BUG=part of tensorflow#3256
Add a test in which a CONCATENATION reads one compressed and one uncompressed constant input. Insertion must decode only the compressed tensor and leave the uncompressed input untouched. The existing tests pair compressed weights with activations, so none of them pin the behavior for an uncompressed constant beside a compressed one. BUG=part of tensorflow#3256
An empty spec list passed to compress() previously returned an unmodified model silently. Fail early with a clear error instead, since an empty spec is almost certainly a mistake. BUG=part of tensorflow#3256
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 #3400, whose branch had fallen far behind main. Rather
than rebase, the tooling was rebuilt fresh on current main,
commit-by-commit, with tf-free tests that build fixtures via
model_editor instead of TensorFlow.
The first commit is the tflm_py_* shim-wrapper change submitted
separately as #3573; the remaining commits are the DECODE tooling.
See the linked issue for a description of the change.
BUG=implements #3256