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

tlopex merged 3 commits into
apache:mainfrom
Aharrypotter:tflite_resource_variable_hashtable

Conversation

@Aharrypotter

@Aharrypotter Aharrypotter commented May 28, 2026 •

Copy link
Copy Markdown
Contributor

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

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

Copy link
Copy Markdown
Contributor

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

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

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
tlopex pushed a commit that referenced this pull request Jun 24, 2026
## 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
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