Skip to content

DECODE fix channel count for non-quantized tensors - #3691

Draft
ddavis-2015 wants to merge 6 commits into
tensorflow:mainfrom
ddavis-2015:decode-channel-fix
Draft

ddavis-2015 wants to merge 6 commits into
tensorflow:mainfrom
ddavis-2015:decode-channel-fix

Conversation

@ddavis-2015

Copy link
Copy Markdown
Member

@tensorflow/micro

Bring DECODE operator LUT and pruning channel count discovery in-line with design documentation. This fixes the channel count always being 1 (one) for non-quantized tensors.

Parse additional field in DCM.

Update unit tests.

bug=fixes #3685

@tensorflow/micro

This is part 1 of 2 to add support for custom DECODE operators and their registration.  Part 2 will add this support to the Python MicroInterpreter wrapper.

Add unit tests for custom DECODE operator registration.

Re-enable -Werror in the Makefile.

Fix minor comment typos.

bug=fixes tensorflow#3215
@tensorflow/micro

Bring DECODE operator LUT and pruning channel count discovery in-line with design documentation. This fixes the channel count always being 1 (one) for non-quantized tensors.
Parse additional field in DCM.
Update unit tests.

bug=fixes tensorflow#3685
@tensorflow/micro

Put the -Werror flag back into the primary makefile, to prevent creaping warning accumulation.

Fix various warning generators in codebase.

bug=fixes tensorflow#3689
@ddavis-2015
ddavis-2015 requested a review from a team as a code owner September 2, 2026 02:20
@ddavis-2015 ddavis-2015 self-assigned this Sep 2, 2026
@ddavis-2015 ddavis-2015 added type:feature New functionality or hardware support implementation. ci:full Triggers the comprehensive cross-platform test suite. labels Sep 2, 2026
@ddavis-2015
ddavis-2015 deployed to integration-test September 2, 2026 02:23 — with GitHub Actions Active
@ddavis-2015
ddavis-2015 marked this pull request as draft September 4, 2026 00:03
@alencheung

Copy link
Copy Markdown

One more release-only edge in the new axis-resolution block, in both the
decode_state_lut.cc and decode_state_prune.cc Setup rewrites:

const uint8_t axis = axis_mask >> kDcmParamsAxisMaskShift;
TFLITE_DCHECK(axis < NumDimensions(&output));
num_channels_ = SizeOfDimension(&output, axis);

axis comes from bits 4-7 of the ancillary params byte — 4 bits of model
data, where the all-ones encoding is the single-channel sentinel, so the
reachable axis values are 0-14 — while the output tensor's rank can be
anything. In release builds the DCHECK is compiled out, and
SizeOfDimension is an unchecked index:

// tensorflow/lite/kernels/kernel_util.h:153-155
inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
  return t->dims->data[dim];
}

So a model with axis >= rank reads past the dims TfLiteIntArray and the
value found there becomes num_channels_. Downstream that value drives:

  • elements_per_channel_ = count_indices_ / num_channels_ (the
    decode_state_lut.cc:56-57 expression, kept by this PR) — a division by
    zero whenever the garbage is 0;
  • the value-table walk span (num_channels_ × channel stride ×
    2^bit_width entries) on the lut side;
  • on the prune side, the zero_points_ allocation size and the
    std::copy_n(quantization->zero_point->data, num_channels_, ...)
    source span.

Since this rework derives the channel count from the shape precisely so
Setup does not depend on the quantization fields agreeing, a release-build
guard in the same shape as the version check at the top of the block would
close it:

TF_LITE_ENSURE_MSG(const_cast<TfLiteContext*>(context_),
                   axis < NumDimensions(&output),
                   "unsupported channel axis %u for output rank %d", axis,
                   NumDimensions(&output));
num_channels_ = SizeOfDimension(&output, axis);

Two adjacent cases in the same block worth covering at the same time:

  1. axis in range but dims[axis] == 0 — re-enters the same
    count_indices_ / num_channels_ divide (also reachable on current main
    via scale->size() == 0). Rejecting num_channels_ == 0 right after
    the assignment turns the crash into an error return:

    TF_LITE_ENSURE_MSG(const_cast<TfLiteContext*>(context_),
                       num_channels_ > 0,
                       "channel axis %u has zero extent", axis);
  2. the agreement checks stay DCHECK-only
    TFLITE_DCHECK(num_channels_ == quantization->scale->size) (lut) and
    == quantization->zero_point->size (prune). With the channel count now
    derived from the shape while the scale/zero-point arrays come from the
    quantization fields, a disagreement between the two sources is exactly
    what a hand-edited model can produce, and in a release build the prune
    side walks zero_point->data with num_channels_ entries
    (std::any_of / std::copy_n) over a shorter array. Promoting both to
    TF_LITE_ENSURE would match the intent of the rework.

The generator side already validates the axis at model-creation time
(compression/lut.py, resolve_mode: axis < rank, axis in (0, rank-1)), but the runtime parse is what a hand-edited model bypasses, so
the guard belongs in Setup rather than only in the tooling.

@ddavis-2015

Copy link
Copy Markdown
Member Author

@alencheung Please do not submit unsolicited comments on draft PRs.

This branch was successfully deployed

1 active deployment
integration-test 40bad14c Deployed Sep 2, 2026 by ddavis-2015 via approval-gate #815
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:full Triggers the comprehensive cross-platform test suite. type:feature New functionality or hardware support implementation.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DECODE fix to number of channels for non-quantized tensors

2 participants