Skip to content

[Relax][Frontend][TFLite] Add TFLite Resource Variable and Static Hashtable Import Support#19639

Merged
tlopex merged 3 commits into
apache:mainfrom
Aharrypotter:tflite_resource_variable_hashtable
May 29, 2026
Merged

[Relax][Frontend][TFLite] Add TFLite Resource Variable and Static Hashtable Import Support#19639
tlopex merged 3 commits into
apache:mainfrom
Aharrypotter:tflite_resource_variable_hashtable

Conversation

@Aharrypotter
Copy link
Copy Markdown
Contributor

@Aharrypotter Aharrypotter commented May 28, 2026

Summary

This PR adds incremental Relax TFLite frontend support for the resource
variable initialization subset:

  • VAR_HANDLE
  • ASSIGN_VARIABLE
  • READ_VARIABLE

It builds on the TFLite control-flow / multi-subgraph support from #19616,
especially CALL_ONCE. TFLite commonly represents initialization through a
CALL_ONCE init subgraph, then uses resource handles from the main subgraph to
read 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:

  • HASHTABLE
  • HASHTABLE_IMPORT
  • HASHTABLE_FIND
  • HASHTABLE_SIZE

These operators are intentionally left unsupported for now. TFLite builtin
hashtable kernels are not generic tensor maps: their runtime implementations
cover the int64 -> string and string -> int64 table variants, and correct
import requires proper TensorType.STRING support. Rejecting the operators is
safer 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_ONCE init
subgraph converter. Each converter instance still keeps its own local
self.resource_handles map, keyed by TFLite tensor name.

Resource variables use container + shared_name from VarHandleOptions when
present, 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_ONCE now accepts a non-empty init subgraph when all operators are in the
supported initialization subset:

  • VAR_HANDLE
  • ASSIGN_VARIABLE

The 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 ExprTable and shared conversion state.

The init subconverter deliberately shares the parent BlockBuilder. This is
safe for the current subset because all supported init operators update importer
state and return None; they do not emit Relax bindings. A comment documents
that this should be revisited if future CALL_ONCE init operators emit Relax
expressions.

Resource Variables

VAR_HANDLE is declarative. It registers the output resource tensor in the
current converter's local resource_handles map and returns None.

ASSIGN_VARIABLE is accepted only while converting a supported CALL_ONCE init
subgraph. 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_VARIABLE resolves the main graph resource handle and returns the
initialized expression from shared state. If the resource has not been
initialized by a supported CALL_ONCE path, the frontend raises
OpNotImplemented.

This supports the common static-initialization inference pattern while avoiding
incorrect lowering for runtime mutation.

Hashtable Operators

HASHTABLE registers the table handle and validates the dtype pair against
TFLite kernel constraints (int64/string or string/int64).

HASHTABLE_IMPORT in a supported CALL_ONCE init subgraph captures static
metadata (table size, key/value dtypes) but does not store actual string data,
because Relax does not yet support TensorType.STRING.

HASHTABLE_SIZE returns a scalar Relax constant for statically imported
tables.

HASHTABLE_FIND is rejected with OpNotImplemented because Relax cannot
represent TFLite string tensors or the runtime lookup semantics.

Operator Support

Operator TFLite options Relax lowering Supported subset
VAR_HANDLE VarHandleOptions handle registration only main graph and supported CALL_ONCE init subgraphs
ASSIGN_VARIABLE AssignVariableOptions store initialized Relax expression in shared importer state supported CALL_ONCE init subgraphs only
READ_VARIABLE ReadVariableOptions return initialized Relax expression resource must have supported static initialization
HASHTABLE HashtableOptions handle registration + dtype validation validates int64/string or string/int64 pair, rejects other combinations
HASHTABLE_IMPORT HashtableImportOptions store static metadata (size, key/value dtype) CALL_ONCE init subgraphs only, constant key/value shape validation
HASHTABLE_FIND HashtableFindOptions unsupported guard requires future TensorType.STRING support in Relax
HASHTABLE_SIZE HashtableSizeOptions scalar Relax constant returns [size] int64 for statically imported tables

Safety Checks

  • ASSIGN_VARIABLE outside CALL_ONCE initialization raises
    OpNotImplemented.
  • READ_VARIABLE without supported initialization raises OpNotImplemented.
  • CALL_ONCE init subgraphs with inputs or outputs remain unsupported.
  • CALL_ONCE init subgraphs containing operators outside the resource-variable
    initialization allowlist remain unsupported.
  • TFLite builtin hashtable operators raise OpNotImplemented until the
    frontend can model their real int64/string table semantics.

Not Included

  • Runtime ASSIGN_VARIABLE mutation in the main graph.
  • Runtime resource-state threading through Relax function parameters and
    returns.
  • Cross-subgraph resource handle aliasing beyond the static
    container/shared_name matching pattern.
  • Multiple runtime writes with ordering semantics.
  • TFLite builtin hashtable lowering.
  • TensorType.STRING import support.

Tests

The tests manually build minimal TFLite flatbuffers and compare imported Relax
IR with tvm.ir.assert_structural_equal. Unsupported patterns use
pytest.raises.

Test Coverage
test_resource_variable_call_once_init_read CALL_ONCE init subgraph with VAR_HANDLE + ASSIGN_VARIABLE, then main graph READ_VARIABLE
test_assign_variable_main_subgraph_unsupported runtime/main graph ASSIGN_VARIABLE guard
test_read_variable_uninitialized_unsupported READ_VARIABLE without supported initialization guard
test_hashtable_call_once_import_find_unsupported hashtable init/find path remains unsupported
test_hashtable_call_once_import_size_unsupported hashtable init/size path remains unsupported
test_hashtable_import_main_subgraph_unsupported main graph HASHTABLE_IMPORT remains unsupported
test_hashtable_size_uninitialized_unsupported uninitialized HASHTABLE_SIZE remains unsupported

Local 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 -q

Result:

py_compile: passed
ruff format --check: files already formatted
ruff check: All checks passed
targeted resource/hashtable tests: 6 passed
full test_frontend_tflite.py: 472 passed

@Aharrypotter Aharrypotter marked this pull request as ready for review May 28, 2026 14:54
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
@Aharrypotter Aharrypotter marked this pull request as draft May 28, 2026 14:56
@Aharrypotter Aharrypotter force-pushed the tflite_resource_variable_hashtable branch from c5f070d to 81aa1be Compare May 29, 2026 03:45
@Aharrypotter Aharrypotter marked this pull request as ready for review May 29, 2026 03:53
@Aharrypotter
Copy link
Copy Markdown
Contributor Author

cc @tlopex

Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@tlopex tlopex changed the title Add TFLite Resource Variable and Static Hashtable Import Support [Relax][Frontend][TFLite] Add TFLite Resource Variable and Static Hashtable Import Support May 29, 2026
@Aharrypotter Aharrypotter marked this pull request as draft May 29, 2026 10:56
@Aharrypotter Aharrypotter force-pushed the tflite_resource_variable_hashtable branch from 8e01581 to 43d5042 Compare May 29, 2026 17:05
@Aharrypotter
Copy link
Copy Markdown
Contributor Author

Hi @tlopex, thanks for the detailed review.

I've removed the incorrect HASHTABLE_FIND lowering and the synthetic INT32 -> INT32 / vector-value tests.

The hashtable handling is now limited to TFLite's actual kernel constraints:

  • HASHTABLE validates table dtype pairs (int64 -> string or string -> int64)
  • HASHTABLE_IMPORT is accepted only in CALL_ONCE init subgraphs, requires matching key/value dtypes and identical key/value shapes, and records only static metadata (size, dtypes)
  • HASHTABLE_SIZE lowers from the recorded metadata
  • HASHTABLE_FIND is rejected with OpNotImplemented because correct lowering requires TensorType.STRING support in Relax

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.

@Aharrypotter Aharrypotter marked this pull request as ready for review May 29, 2026 17:22
@Aharrypotter
Copy link
Copy Markdown
Contributor Author

Update PR description

@Aharrypotter Aharrypotter requested a review from tlopex May 29, 2026 17:28
@tlopex tlopex merged commit b971a75 into apache:main May 29, 2026
11 of 12 checks passed
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.

2 participants