feat(compression): add DECODE operator insertion by rkuester · Pull Request #3624 · tensorflow/tflite-micro

rkuester · 2026-07-06T20:42:36Z

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.

Compression tooling only; adds decode_insert.py and its test, no
runtime changes.

Part of the DECODE compression-tooling stack.

BUG=#3256

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

veblush · 2026-07-06T23:18:04Z

+    if not consumers:
+      # Check if tensor is a subgraph output
+      is_output = tensor in subgraph.outputs
+      if is_output:


I guess if it_output should be outside of if not consumers? It might be consumed and an oiutput of subgraph?

You're right: the check missed tensors that are both consumed by an operator and listed as a subgraph output, so a subgraph would have output the encoded bytes instead of the decoded values. Addressed in 3999889 by treating the output list as one more consumer, fed by a DECODE appended after the last operator.

f3a8e56 pulls forward a fix from later in the series: separate DECODEs per tensor don't work with alternate decompression memory when the decoded values are needed at the same time, and with output handling added, it's even more worth getting this right initially than rewriting it later. It's a separate commit for easy review, relying on the squash when the PR merges.

@rkuester @veblush Re: Passing the output of DECODE to or from a subgraph

The design document will be updated to address what to do in this situation.

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: 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: consumers are handled in reverse position order, so each insertion falls after every consumer still to be processed. BUG=part of tensorflow#3256

ddavis-2015

Looks really good, almost there. What is here will create a valid and usable model, but not exactly as it should be.

Maximal reduction of the model is not yet achieved due to possible duplication of the ancillary tensor Buffer. Here is the scenario: Two const-tensors, either in the same or in different subgraphs can contain the exact same data. The TfLite Converter can (and does) create two tensors in the model, but using a single shared buffer. For DECODE insertion, the existing code works just fine for encoded tensors, as they are not given a new buffer, just new buffer content. But this means that the related ancillary tensors, can also share a single buffer. However, the current insertion code always creates a new Buffer object when creating ancillary tensors.

The output tensors of DECODE must be exact copies of the original unencoded tensors, meaning when the output tensors are created they must use the original tensor's TensorT object. Only the name of the output tensors can be modified (as is currently done by the insertion code), and the Buffer object set to None (buffer in the TensorT object is zero).

For the unit tests, maybe a comment could be added stating the models are just for unit test purposes and do not represent fully valid models. Same for the encoded_data and ancillary_data of the CompressionResult objects. Something saying these are just dummy payloads for unit testing and do not correspond to the model's tensor data. Someone looking at these for the first time might mistake these for valid examples and be confused.

ddavis-2015 · 2026-07-13T17:56:12Z

+    sg = model.subgraphs[0]
+    weights = sg.tensor_by_name("shared_weights")
+
+    consumers = decode_insert._find_tensor_consumers(sg, weights)


Use of a private method (_find_tensor_consumers). Maybe have a set of public methods just for use by unit tests?

ddavis-2015 · 2026-07-13T19:48:31Z

+  model = model_editor.Model(subgraphs=[
+      model_editor.Subgraph(
+          tensors=[weights],
+          operators=[
+              model_editor.Operator(
+                  opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
+                  inputs=[input_t, weights],
+                  outputs=[output_t],
+              )
+          ],
+      )


Doesn't the Subgraph inputs/outputs need to be specified? Otherwise the subgraph compiler will just assign empty vectors to the tflite.SubGraphT.

ddavis-2015 · 2026-07-13T19:52:40Z

+  model = model_editor.Model(subgraphs=[
+      model_editor.Subgraph(
+          tensors=[table],
+          operators=[
+              model_editor.Operator(
+                  opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
+                  inputs=[input_t, weights],
+                  outputs=[output_t],
+              )
+          ],
+          outputs=[output_t, table],
+      )
+  ])


Doesn't the Subgraph inputs need to be specified? Otherwise the subgraph compiler will just assign an empty vector in the tflite.SubGraphT.

ddavis-2015 · 2026-07-13T20:15:01Z

+            encoded_data=b'\x00\x00',
+            ancillary_data=_make_dummy_ancillary_data(),


mismatch between encoded_data and the ancillary_data. Perhaps _make_dummy_ancillary_data should have arguments so that it can generate a DCM that matches the model data? And encoded_data appears to be too short in length.

ddavis-2015 · 2026-07-13T20:21:03Z

+  def test_insert_single_decode_operator(self):
+    """DECODE operator inserted before FC that uses compressed weights."""
+    model = _build_simple_fc_model()
+    weights_tensor = model.subgraphs[0].tensor_by_name("weights")


unused variable weights_tensor

ddavis-2015 · 2026-07-15T03:37:12Z

+  model = model_editor.Model(subgraphs=[
+      model_editor.Subgraph(
+          tensors=[weights],
+          operators=[
+              model_editor.Operator(
+                  opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
+                  inputs=[input1, weights],
+                  outputs=[output1],
+              ),
+              model_editor.Operator(
+                  opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
+                  inputs=[input2, weights],
+                  outputs=[output2],
+              ),
+          ],
+      )


Doesn't the Subgraph inputs/outputs need to be specified? Otherwise the subgraph compiler will just assign empty vectors to the tflite.SubGraphT.

ddavis-2015 · 2026-07-15T04:03:21Z

+    # The two DECODEs should share the same ancillary tensor
+    self.assertIs(decode_op1.inputs[1], decode_op2.inputs[1])


also check the two DECODES share the same encoded tensor

ddavis-2015 · 2026-07-15T04:30:23Z

+    self.assertIs(consumer_decode.inputs[1], output_decode.inputs[1])
+


also check both DECODEs share the encoded tensor

ddavis-2015 · 2026-07-15T22:58:00Z

+                    opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
+                    inputs=[input_t, weights1, weights2],


CONCATENATION would be a better example here (with output shape [6,4])

ddavis-2015 · 2026-07-16T01:26:25Z

+
+    # Verify DCM header
+    dcm_bytes = ancillary_tensor.array[:16]
+    self.assertEqual(dcm_bytes[0], 0)  # decode_type = LUT


use DecodeType.LUT

rkuester requested a review from a team as a code owner July 6, 2026 20:42

rkuester added the ci:full Triggers the comprehensive cross-platform test suite. label Jul 6, 2026

rkuester temporarily deployed to integration-test July 6, 2026 20:42 — with GitHub Actions Inactive

rkuester requested a review from veblush July 6, 2026 22:34

veblush reviewed Jul 6, 2026

View reviewed changes

Comment thread tensorflow/lite/micro/compression/decode_insert.py Outdated

rkuester added 3 commits July 7, 2026 12:50

rkuester temporarily deployed to integration-test July 7, 2026 18:48 — with GitHub Actions Inactive

rkuester requested review from ddavis-2015 and veblush July 7, 2026 19:04

ddavis-2015 suggested changes Jul 16, 2026

View reviewed changes

rkuester mentioned this pull request Jul 18, 2026

WHILE operator input/output copy fix #3633

Merged

		encoded_data=b'\x00\x00',
		ancillary_data=_make_dummy_ancillary_data(),

		# The two DECODEs should share the same ancillary tensor
		self.assertIs(decode_op1.inputs[1], decode_op2.inputs[1])

		self.assertIs(consumer_decode.inputs[1], output_decode.inputs[1])

		opcode=tflite.BuiltinOperator.FULLY_CONNECTED,
		inputs=[input_t, weights1, weights2],

Uh oh!

Conversation

rkuester commented Jul 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ddavis-2015 Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ddavis-2015 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ddavis-2015 Jul 8, 2026 •

edited

Loading

ddavis-2015 left a comment •

edited

Loading