Skip to content

[Relax][Frontend][TFLite] Read the keys vector byte width in FlexBufferDecoder - #20830

Open
mmporong wants to merge 1 commit into
apache:mainfrom
mmporong:fix-tflite-flexbuffer-key-width
Open

mmporong wants to merge 1 commit into
apache:mainfrom
mmporong:fix-tflite-flexbuffer-key-width

Conversation

@mmporong

Copy link
Copy Markdown

FlexBufferDecoder.decode_map decoded the keys vector of a map with 1-byte offsets no matter what the buffer said. A FlexBuffers map header stores the byte width of its keys vector right after the keys offset. The FlexBuffers builder widens that vector to 2 bytes once a key string sits more than 255 bytes before it, which happens with long keys or a few dozen short ones. For such a map the decoder read wrong key offsets and returned garbage keys, or raised UnicodeDecodeError.

This change reads the keys byte width from the map header and uses it for the key offsets. The custom options of TFLite_Detection_PostProcess, the only current caller, keep a 1-byte keys vector, so they decode exactly as before. That was checked with the standard 11 options and with _output_quantized and _support_output_type_float_in_quantized_op added.

The new test_flexbuffer_decoder_wide_key_vector in tests/python/relax/test_frontend_tflite.py encodes two flat maps with flatbuffers.flexbuffers, one with long keys and one with 40 keys, and checks that they decode back to the original dict. libtvm was not built locally, so the test body was run as is, with tflite_flexbuffer.py loaded by path (the module only needs the standard library). Without the fix both cases fail, one with UnicodeDecodeError and one with wrong keys. With the fix both pass. A wider comparison against flexbuffers.Loads also matches, covering maps of up to 200 entries and mixed int, float and bool values. pre-commit run --files on both files passes (ruff check and ruff format). The repository-wide ASF header and file type hooks were skipped in the sparse checkout used here.

…erDecoder

FlexBufferDecoder.decode_map always decoded the keys vector of a map with
1-byte offsets. A FlexBuffers map header stores the byte width of its keys
vector next to the keys offset, and the builder widens it to 2 bytes once
a key string is more than 255 bytes before the vector, for example with
long or many keys. Such a map was then decoded into wrong keys or raised
UnicodeDecodeError.

Read the keys byte width from the map header and use it for the key
offsets. The custom options of TFLite_Detection_PostProcess, the only
current caller, keep a 1-byte keys vector and decode as before.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant