LiteRT / TFLite Export to ONNX#

yobx.litert.to_onnx() converts a TFLite/LiteRT .tflite model into an onnx.ModelProto. The implementation is a proof-of-concept that parses the binary FlatBuffer format of .tflite files with a minimal pure-Python parser (no external library required) and converts each operator in the graph to its ONNX equivalent via a registry of op-level converters.

High-level workflow#

.tflite file / bytes
       │
       ▼
   to_onnx()           ← parse_tflite_model() + build TFLiteModel
       │
       ▼
TFLiteModel             (pure-Python object hierarchy)
       │
       ▼
_convert_subgraph()
       │  for every operator in the subgraph …
       ▼
op converter            ← emits ONNX node(s) via GraphBuilder.op.*
       │
       ▼
GraphBuilder.to_onnx()  ← validates and returns ModelProto

The steps in detail:

  1. to_onnx accepts the model as a file path or raw bytes together with a tuple of representative numpy inputs (used to determine input dtypes and shapes when they are not fully specified in the TFLite model).

  2. parse_tflite_model() reads the binary FlatBuffer and returns a TFLiteModel object containing a list of TFLiteSubgraph objects, each with its tensors and operators.

  3. A fresh GraphBuilder is created and _convert_subgraph() walks the operator list:

    1. Weight tensors (tensors that carry buffer data and are not graph inputs) are registered as ONNX initializers.

    2. Input tensors are registered via make_tensor_input().

    3. Every operator is dispatched to a registered converter (or to an entry in extra_converters).

  4. The ONNX outputs are declared with make_tensor_output().

  5. GraphBuilder.to_onnx finalises and returns the onnx.ModelProto.

Quick example#

import numpy as np
from yobx.litert import to_onnx

X = np.random.rand(1, 4).astype(np.float32)
onx = to_onnx("model.tflite", (X,))

TFLite FlatBuffer parser#

yobx.litert.litert_helper contains a self-contained, zero-dependency FlatBuffer reader (_FlatBuf) together with the parsed data-classes (TFLiteModel, TFLiteSubgraph, TFLiteTensor, TFLiteOperator) and the BuiltinOperator enum.

<<<

from yobx.litert.litert_helper import (
    _make_sample_tflite_model,
    parse_tflite_model,
)

model = parse_tflite_model(_make_sample_tflite_model())
sg = model.subgraphs[0]
print(f"subgraphs : {len(model.subgraphs)}")
print(f"tensors   : {[t.name for t in sg.tensors]}")
print(f"inputs    : {sg.inputs}")
print(f"outputs   : {sg.outputs}")
for op in sg.operators:
    print(f"operator  : {op.name}  inputs={op.inputs}  outputs={op.outputs}")

>>>

    subgraphs : 1
    tensors   : ['input', 'relu']
    inputs    : (0,)
    outputs   : (1,)
    operator  : RELU  inputs=(0,)  outputs=(1,)

Converter registry#

The registry is a module-level dictionary LITERT_OP_CONVERTERS: Dict[Union[int, str], Callable] defined in yobx.litert.register. Keys are BuiltinOperator integers (e.g. BuiltinOperator.RELU = 19); custom ops use their string name. Values are converter callables.

Registering a converter#

Use the register_litert_op_converter() decorator. Pass a single op-code integer, a custom-op name string, or a tuple thereof:

from yobx.litert.register import register_litert_op_converter
from yobx.litert.litert_helper import BuiltinOperator


@register_litert_op_converter(BuiltinOperator.RELU)
def convert_relu(g, sts, outputs, op):
    return g.op.Relu(op.inputs[0], outputs=outputs, name="relu")

Converter function signature#

Every op converter follows the same contract:

(g, sts, outputs, op) → output_name

Parameter

Description

g

GraphBuilder — call g.op.<OpType>(…) to emit ONNX nodes.

sts

Dict of metadata (currently always {}).

outputs

List[str] of pre-allocated output tensor names that the converter must write to.

op

An _OpProxy whose inputs and outputs are string names of the tensors, and builtin_options is a decoded attribute dict.

op.inputs[i] is the ONNX tensor name (a string, not an integer index) of the i-th operator input. Use it directly in g.op.*() calls.

Dynamic shapes#

By default to_onnx marks axis 0 of every input as dynamic (unnamed batch dimension). To control which axes are dynamic, pass dynamic_shapes — a tuple of one Dict[int, str] per input where keys are axis indices and values are symbolic dimension names:

onx = to_onnx("model.tflite", (X,), dynamic_shapes=({0: "batch"},))

Custom op converters#

The extra_converters parameter of to_onnx accepts a mapping from BuiltinOperator integer (or custom-op name string) to converter function. Entries here take priority over the built-in registry:

import numpy as np
from yobx.litert import to_onnx
from yobx.litert.litert_helper import BuiltinOperator

def custom_relu(g, sts, outputs, op):
    """Replace RELU with Clip(0, 6)."""
    return g.op.Clip(
        op.inputs[0],
        np.array(0.0, dtype=np.float32),
        np.array(6.0, dtype=np.float32),
        outputs=outputs,
        name="relu6",
    )

onx = to_onnx("model.tflite", (X,),
              extra_converters={BuiltinOperator.RELU: custom_relu})

Adding a new built-in converter#

  1. Create a new file under yobx/litert/ops/ (e.g. yobx/litert/ops/cast_ops.py).

  2. Implement a converter function following the signature above.

  3. Decorate it with @register_litert_op_converter(BuiltinOperator.CAST).

  4. Import the new module inside the register() function in yobx/litert/ops/__init__.py.

# yobx/litert/ops/cast_ops.py
from onnx import TensorProto
from ..register import register_litert_op_converter
from ..litert_helper import BuiltinOperator
from ...xbuilder import GraphBuilder

@register_litert_op_converter(BuiltinOperator.CAST)
def convert_cast(g: GraphBuilder, sts: dict, outputs: list, op) -> str:
    """TFLite CAST → ONNX Cast."""
    # op.builtin_options carries "in_data_type" / "out_data_type" integers.
    out_type = op.builtin_options.get("out_data_type", 0)
    from yobx.litert.litert_helper import litert_dtype_to_np_dtype
    import numpy as np
    np_dtype = litert_dtype_to_np_dtype(out_type)
    onnx_dtype = TensorProto.FLOAT  # map np_dtype → TensorProto int
    return g.op.Cast(op.inputs[0], to=onnx_dtype, outputs=outputs, name="cast")

Supported ops#

See Supported LiteRT Ops for the full list of built-in LiteRT op converters, generated automatically from the live registry.