Skip to content

fuse_rolled_gru

Fuse the canonical Keras rolled-GRU loop into one native GRU operator.

This runs at the LiteRT level rather than as an AIR transform because AIR has no representation for WHILE or for multi-subgraph models: an AIR-level pass could never see this graph, since parsing would fail before it ran. The pass is registered as a default :data:~helia_aot.litert.passes.LiteRtModelHook so it stays ordered, opt-out-able, and overridable through RegistryContext.

The matcher verifies the GRU equations by walking the cell body rather than keying on tensor names, because the fully-FP16 models this path requires cannot come from the stock TensorFlow converter -- they are customer rewrites of it, which is exactly the population where incidental export artifacts vary.

Copyright 2025 Ambiq. All Rights Reserved.

Classes

Functions

fuse_rolled_gru_ops 

fuse_rolled_gru_ops(model: ModelT, *, subgraph_index: int = 0) -> int

Replace recognized Keras WHILE GRUs with canonical CUSTOM(GRU).

The input LiteRT model remains an ordinary TensorFlow-converted rolled GRU. This pass is an internal lowering step performed immediately before AIR parsing.

Parameters:

  • model

    (ModelT) –

    Mutable LiteRT ObjectAPI model.

  • subgraph_index

    (int, default: 0 ) –

    Main subgraph containing the rolled GRU.

Returns:

  • int –

    Number of loops fused.

Raises:

  • UnsupportedModelError –

    If a loop is verified to be a Keras GRU cell but is not fully FP16, has non-constant projection parameters, or the model already declares a colliding CUSTOM("GRU") opcode.

fuse_rolled_gru_hook 

fuse_rolled_gru_hook(model: ModelT, subgraph_index: int, model_path: Path) -> litert.ModelT | None

Default :data:LiteRtModelHook wrapper for :func:fuse_rolled_gru_ops.

Parameters:

  • model

    (ModelT) –

    Mutable LiteRT ObjectAPI model.

  • subgraph_index

    (int) –

    Main subgraph containing the rolled GRU.

  • model_path

    (Path) –

    Source model path, unused.

Returns:

  • ModelT | None –

    None, since the model is rewritten in place.