From e3c7f75665d1c3dcb9a6be0130c384b6bea19f9d Mon Sep 17 00:00:00 2001
From: mmporong <mmporong@gmail.com>
Date: Tue, 27 Oct 2026 16:19:06 +0900
Subject: [PATCH] [Relax][Frontend][TFLite] Read the keys vector byte width in
 FlexBufferDecoder

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.
---
 .../relax/frontend/tflite/tflite_flexbuffer.py  |  9 +++++++--
 tests/python/relax/test_frontend_tflite.py      | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relax/frontend/tflite/tflite_flexbuffer.py b/python/tvm/relax/frontend/tflite/tflite_flexbuffer.py
index 148f404f25db..ce415fdece9c 100644
--- a/python/tvm/relax/frontend/tflite/tflite_flexbuffer.py
+++ b/python/tvm/relax/frontend/tflite/tflite_flexbuffer.py
@@ -135,10 +135,15 @@ def decode_map(self, end, byte_width, parent_byte_width):
         size_fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[byte_width]
         map_size = struct.unpack(size_fmt, self.buffer[mid_loc - byte_width : mid_loc])[0]
 
-        # Find keys
+        # Find keys. The map header also stores the byte width of the keys vector, which is
+        # wider than one byte once a key string sits more than 255 bytes before it.
         keys_offset = mid_loc - byte_width * 3
         keys_end = self.indirect_jump(keys_offset, byte_width)
-        keys = self.decode_keys(keys_end, map_size, 1)
+        keys_byte_width_fmt = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[byte_width]
+        keys_byte_width = struct.unpack(
+            keys_byte_width_fmt, self.buffer[mid_loc - byte_width * 2 : mid_loc - byte_width]
+        )[0]
+        keys = self.decode_keys(keys_end, map_size, keys_byte_width)
 
         # Find values
         values_end = self.indirect_jump(end, parent_byte_width)
diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py
index 4799c3481cfd..db67e2c26b7b 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -3741,6 +3741,23 @@ def test_detection_postprocess_rejects_invalid_class_count():
         _build_detection_postprocess_mod(num_classes=2, input_num_classes=5)
 
 
+@pytest.mark.parametrize(
+    "custom_options",
+    [
+        pytest.param({f"option_{i}_" + "x" * 40: i for i in range(8)}, id="long_keys"),
+        pytest.param({f"key_{i:03d}": i for i in range(40)}, id="many_keys"),
+    ],
+)
+def test_flexbuffer_decoder_wide_key_vector(custom_options):
+    # Over 255 bytes of key strings make the keys vector use 2-byte offsets.
+    from flatbuffers import flexbuffers
+
+    from tvm.relax.frontend.tflite.tflite_flexbuffer import FlexBufferDecoder
+
+    flexbuffer = bytes(flexbuffers.Dumps(custom_options))
+    assert FlexBufferDecoder(flexbuffer).decode() == custom_options
+
+
 def _make_resize_expected(
     input_shape, output_size, method, coordinate_transformation_mode, rounding_method
 ):