Stateful models

A recurrent model carries state in one of two ways. Find out which one the exported model uses before writing its streaming loop or its golden fixture.

Operator or path	State contract
LSTM and SVDF internal state	Persistent variable tensors survive aot_model_run; aot_model_init resets them. Check each operator’s dtype and shape restrictions in the catalog.
Supported rolled GRU	Explicit initial-state input and final-state output; the caller carries the state. Persistent GRU state tensors are rejected.
Resource variables	READ_VARIABLE and ASSIGN_VARIABLE use persistent storage according to the graph’s variable contract.

The catalog rows give the data types: LSTM and SVDF have int8 and float paths with role-specific conditions; GRU is FP16, batch size one, with reset_after=True. Read the whole row for weights, sequence layout, state and rank constraints.

Where the state lives

Internal variable state is whole-program-live writable storage in a persistent arena. aot_model_init resets the persistent slots and runs operator initialization; aot_model_run keeps the resulting state for the next call.

1. Initialize once for a new independent sequence and check the status.
2. Copy the next input window and run; check the status before consuming the output.
3. Repeat input and run for the following windows without reinitializing.
4. Initialize again to reset the sequence on purpose.

For explicitly exposed state, copy the previous state output to the matching state input, using the model’s I/O order and matching shape and dtype. Use memmove when the buffers may overlap. Reset by supplying the intended initial-state values before the next independent sequence.

Bindings and generated arena buffers are module-global. Serialize calls and state updates; two context structs are not independent stateful instances. The residency report shows persistent storage, while explicit GRU state follows the model’s I/O contract. Memory covers arena ownership and sharing.

The rolled GRU

A Keras export can represent the GRU as a WHILE loop across LiteRT subgraphs. The default fuse_rolled_gru LiteRT model hook recognizes the supported cell equations and fuses that form before conversion to AIR. It is distinct from the AIR graph transforms applied later.

The supported GRU has sequence and initial-state inputs, plus sequence and final-state outputs. State tensors must be ordinary graph I/O, not persistent LiteRT variables. Check the model-level I/O order: it need not match the operator’s local tensor order.

Python registry customizers can replace or remove the hook, but removing it does not make the raw loop lowerable. A different export, or a custom parser and lowering, is then required.

Carrying state through a test run

For internal persistent state, repeated test iterations keep the state without test.state_feedback. Non-state inputs are refreshed from the same stimulus on each iteration; this is not a stream of different input windows.

For explicit state I/O, test.state_feedback lists zero-based [output_position, input_position] pairs. Positions refer to the model’s output and input lists, not flatbuffer tensor IDs. For a model whose final state is output 0 and initial state is input 1:

test:
  enabled: true
  golden_data: golden.npz
  num_iterations: 4
  state_feedback:
    - [0, 1]

Use this mapping only after checking your model’s order. The converter rejects out-of-range positions, shape or dtype mismatches, and two outputs assigned to one input. The harness seeds the state input on iteration zero, then carries the output into it with memmove; the other inputs are refreshed each iteration.

With a supplied NPZ, output_N must describe the final iteration under that same stimulus, initial state, count and carry mapping. The converter does not rerun supplied golden data to repair a mismatched protocol. Without supplied golden data, its host reference loop uses the configured feedback and count.

For a real sequence of different windows, use an application harness that feeds those windows and checks the required outputs. The generated repeated-stimulus test alone does not establish streaming accuracy. See Testing and automation for the test contract.