Repository navigation
[Relax][Frontend][TFLite] Add TFLite Resource Variable and Static Hashtable Import Support - #19639
Conversation
There was a problem hiding this comment.
Code Review
This pull request adds support for resource variables and static hashtables in the Relax TFLite frontend. It introduces tracking for resource and hashtable values, supports several new operators (such as VAR_HANDLE, ASSIGN_VARIABLE, READ_VARIABLE, HASHTABLE, HASHTABLE_IMPORT, HASHTABLE_FIND, and HASHTABLE_SIZE), and updates CALL_ONCE handling to process non-empty initialization subgraphs. Comprehensive unit tests are also added. The review feedback highlights several robust improvements: wrapping optional TFLite imports in try-except blocks to prevent import errors on older environments, explicitly validating that hashtable keys and values are constant tensors, and correctly handling 0D (scalar) query keys in HASHTABLE_FIND.
c5f070d to
81aa1be
Compare
|
cc @tlopex |
tlopex
left a comment
There was a problem hiding this comment.
Thanks for working on this. I think the hashtable part does not currently match TFLite's builtin HASHTABLE semantics.
The TFLite kernels only support int64 -> string or string -> int64 tables. HASHTABLE_IMPORT also requires the key and value tensors to have the same shape, and HASHTABLE_FIND resizes the output to the query key tensor shape.
This PR's positive tests build INT32 -> INT32 tables and the converter supports imported values with shape [num_keys, ...], returning output shaped like query_shape + value_shape. That is not the contract accepted by the TFLite runtime. Conversely, actual valid TFLite hashtable models involving STRING tensors cannot be imported here because the frontend tensor dtype helpers do not support TensorType.STRING.
So this can accept synthetic/malformed hashtable flatbuffers while not supporting the real TFLite hashtable subset it claims to add. I think the converter should either match TFLite's kernel constraints and add positive tests for valid int64/string and string/int64 tables, or reject HASHTABLE_* until those semantics are supported. The current INT32 and vector-value tests should not be the positive coverage for this operator family.
8e01581 to
43d5042
Compare
|
Hi @tlopex, thanks for the detailed review. I've removed the incorrect The hashtable handling is now limited to TFLite's actual kernel constraints:
So this PR no longer models hashtables as a generic tensor map. It keeps the resource-variable support plus the safe hashtable metadata/size subset, and leaves real string lookup lowering for a future PR after string tensor support exists. Please take another look when you have time. |
|
Update PR description |
## Summary This PR extends the TFLite resource hashtable family from #19519 item G from "import + size" to a constant-foldable `HASHTABLE_FIND` subset. It builds on the static `HASHTABLE` / `HASHTABLE_IMPORT` import support added in #19639 and the `HASHTABLE_LOOKUP` converter from #19654. Relax has no string tensor type and no runtime hashtable operator, so this PR targets the only subset that can be lowered today: tables whose keys and values are constants imported in a `CALL_ONCE` init subgraph, queried by a constant string tensor. In that case the lookup is resolved at import time and emitted as a `relax.const`. Runtime string queries remain explicitly guarded instead of being lowered with incorrect semantics. ## Design ### Static String Tensor Decoding The frontend now decodes constant TFLite string tensors through a shared helper `_get_string_tensor_value`. It parses the TFLite string-tensor binary layout: ```text [count(i32)][offset_0 .. offset_count(i32)][utf8 string data] ``` The helper validates the buffer length against the declared count, checks that offsets are in-bounds and monotonically non-decreasing, and verifies the decoded element count matches the tensor shape. This is reusable infrastructure for any future TFLite string support, independent of `HASHTABLE_FIND`. ### Hashtable Import State `HASHTABLE_IMPORT` now stores the actual constant keys and values (numeric or decoded string buffers) in shared conversion state, instead of only recording table metadata (size, key/value dtype). Duplicate keys are rejected, because the constant-fold lookup assumes a unique key mapping. The captured table is keyed by the same `table_id` / handle resolution used by `HASHTABLE` and `HASHTABLE_SIZE`, so a `CALL_ONCE` init subgraph and the main graph agree on the same logical table. ### Constant-Foldable Find `HASHTABLE_FIND` resolves the table handle through the importer-local handle map and the statically imported keys/values. For the supported subset it builds a Python key -> value map, applies the per-element default value, overwrites hits from the table, preserves the query tensor shape, and emits the result as a `relax.const`. The op produces no runtime Relax computation, which matches the fact that both the table and the query are compile-time constants. The supported subset is intentionally narrow and guard-first: - the table must be a constant `string -> int64` table imported via a supported `CALL_ONCE` `HASHTABLE_IMPORT` - the query tensor must be a constant string buffer - the default tensor must be a scalar or match the query shape ### String Graph Input Guard `TensorType.STRING` graph inputs are now rejected in `from_tflite` with a clear `OpNotImplemented` instead of a low-level FFI `unknown dtype string` error, since Relax cannot represent a string tensor. This is the path hit by runtime string queries for `HASHTABLE_FIND` and string-valued `HASHTABLE_LOOKUP`, so the guard gives a frontend-level diagnostic for both. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `HASHTABLE_IMPORT` | `HashtableImportOptions` | store constant keys/values in importer state | `CALL_ONCE` init, constant keys/values, no duplicate keys | | `HASHTABLE_FIND` | `HashtableFindOptions` | constant-fold to `relax.const` | constant `string -> int64` table + constant string query | ## Not Included - Runtime (non-constant) string queries and runtime hashtable lookup. - `int64 -> string` find, which would require string-typed Relax outputs. - Any general `TensorType.STRING` tensor representation in Relax. - A runtime Relax hashtable operator with string hashing / comparison. - Mutable runtime resource-state threading through Relax functions. These require core Relax support and are out of scope for the frontend. ## Tests The tests manually build minimal TFLite flatbuffers and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported patterns use `pytest.raises`. | Test | Coverage | |---|---| | `test_hashtable_call_once_import_find_string_to_int64` | constant `string -> int64` find folds to a `relax.const` | | `test_hashtable_call_once_import_find_string_to_int64_2d_query` | query shape preserved for a 2-D query | | `test_hashtable_call_once_import_find_int64_to_string_unsupported` | `int64 -> string` table rejected | | `test_hashtable_call_once_import_find_runtime_query_unsupported` | runtime string query rejected | | `test_hashtable_call_once_import_duplicate_keys_unsupported` | duplicate static keys rejected | | `test_hashtable_lookup_string_value_unsupported` | string graph input now gives a clean `OpNotImplemented` | Local validation: ```bash python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k "hashtable or resource or variable" -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text ruff format --check: 2 files already formatted ruff check: All checks passed 14 passed, 535 deselected 549 passed ``` ## References - Issue #19519 item G: TFLite resource / variable / hashtable operators - PR #19639: TFLite resource variable and static hashtable import support - PR #19654: TFLite `HASHTABLE_LOOKUP` converter
Summary
This PR adds incremental Relax TFLite frontend support for the resource
variable initialization subset:
VAR_HANDLEASSIGN_VARIABLEREAD_VARIABLEIt builds on the TFLite control-flow / multi-subgraph support from #19616,
especially
CALL_ONCE. TFLite commonly represents initialization through aCALL_ONCEinit subgraph, then uses resource handles from the main subgraph toread initialized variables. This PR supports that constrained initialization
pattern without introducing general mutable runtime state into Relax.
The PR also adds explicit frontend guards for the TFLite builtin hashtable
operators:
HASHTABLEHASHTABLE_IMPORTHASHTABLE_FINDHASHTABLE_SIZEThese operators are intentionally left unsupported for now. TFLite builtin
hashtable kernels are not generic tensor maps: their runtime implementations
cover the
int64 -> stringandstring -> int64table variants, and correctimport requires proper
TensorType.STRINGsupport. Rejecting the operators issafer than lowering a synthetic numeric table semantics that TFLite does not
actually implement.
Design
Shared Initialization State
The frontend now keeps resource initialization data in shared conversion state:
conversion_state["resource_values"]conversion_state["in_call_once_init"]This state is shared by the main graph converter and the
CALL_ONCEinitsubgraph converter. Each converter instance still keeps its own local
self.resource_handlesmap, keyed by TFLite tensor name.Resource variables use
container + shared_namefromVarHandleOptionswhenpresent, falling back to the handle tensor name. This keeps tensor-name bindings
scoped to each subgraph while allowing init subgraphs and the main graph to
agree on the same logical resource.
CALL_ONCE Init Subgraphs
CALL_ONCEnow accepts a non-empty init subgraph when all operators are in thesupported initialization subset:
VAR_HANDLEASSIGN_VARIABLEThe init subgraph still must have no inputs and no outputs. The converter first
checks every operator against the allowlist, then converts the init subgraph
with a fresh
ExprTableand shared conversion state.The init subconverter deliberately shares the parent
BlockBuilder. This issafe for the current subset because all supported init operators update importer
state and return
None; they do not emit Relax bindings. A comment documentsthat this should be revisited if future
CALL_ONCEinit operators emit Relaxexpressions.
Resource Variables
VAR_HANDLEis declarative. It registers the output resource tensor in thecurrent converter's local
resource_handlesmap and returnsNone.ASSIGN_VARIABLEis accepted only while converting a supportedCALL_ONCEinitsubgraph. It resolves the resource handle through the init converter's local
handle map and stores the assigned tensor expression in shared
conversion_state["resource_values"].READ_VARIABLEresolves the main graph resource handle and returns theinitialized expression from shared state. If the resource has not been
initialized by a supported
CALL_ONCEpath, the frontend raisesOpNotImplemented.This supports the common static-initialization inference pattern while avoiding
incorrect lowering for runtime mutation.
Hashtable Operators
HASHTABLEregisters the table handle and validates the dtype pair againstTFLite kernel constraints (
int64/stringorstring/int64).HASHTABLE_IMPORTin a supportedCALL_ONCEinit subgraph captures staticmetadata (table size, key/value dtypes) but does not store actual string data,
because Relax does not yet support
TensorType.STRING.HASHTABLE_SIZEreturns a scalar Relax constant for statically importedtables.
HASHTABLE_FINDis rejected withOpNotImplementedbecause Relax cannotrepresent TFLite string tensors or the runtime lookup semantics.
Operator Support
VAR_HANDLEVarHandleOptionsCALL_ONCEinit subgraphsASSIGN_VARIABLEAssignVariableOptionsCALL_ONCEinit subgraphs onlyREAD_VARIABLEReadVariableOptionsHASHTABLEHashtableOptionsint64/stringorstring/int64pair, rejects other combinationsHASHTABLE_IMPORTHashtableImportOptionsCALL_ONCEinit subgraphs only, constant key/value shape validationHASHTABLE_FINDHashtableFindOptionsTensorType.STRINGsupport in RelaxHASHTABLE_SIZEHashtableSizeOptions[size]int64 for statically imported tablesSafety Checks
ASSIGN_VARIABLEoutsideCALL_ONCEinitialization raisesOpNotImplemented.READ_VARIABLEwithout supported initialization raisesOpNotImplemented.CALL_ONCEinit subgraphs with inputs or outputs remain unsupported.CALL_ONCEinit subgraphs containing operators outside the resource-variableinitialization allowlist remain unsupported.
OpNotImplementeduntil thefrontend can model their real int64/string table semantics.
Not Included
ASSIGN_VARIABLEmutation in the main graph.returns.
container/shared_namematching pattern.TensorType.STRINGimport support.Tests
The tests manually build minimal TFLite flatbuffers and compare imported Relax
IR with
tvm.ir.assert_structural_equal. Unsupported patterns usepytest.raises.test_resource_variable_call_once_init_readCALL_ONCEinit subgraph withVAR_HANDLE + ASSIGN_VARIABLE, then main graphREAD_VARIABLEtest_assign_variable_main_subgraph_unsupportedASSIGN_VARIABLEguardtest_read_variable_uninitialized_unsupportedREAD_VARIABLEwithout supported initialization guardtest_hashtable_call_once_import_find_unsupportedtest_hashtable_call_once_import_size_unsupportedtest_hashtable_import_main_subgraph_unsupportedHASHTABLE_IMPORTremains unsupportedtest_hashtable_size_uninitialized_unsupportedHASHTABLE_SIZEremains unsupportedLocal validation:
python -m py_compile \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k "resource_variable or read_variable_uninitialized or hashtable" -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -qResult: