yobx.litert.ops.nn_ops#

Converters for TFLite neural network ops: FULLY_CONNECTED, CONV_2D, DEPTHWISE_CONV_2D, AVERAGE_POOL_2D, MAX_POOL_2D, BATCH_MATMUL.

yobx.litert.ops.nn_ops.convert_avg_pool2d(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite AVERAGE_POOL_2D → ONNX AveragePool.

yobx.litert.ops.nn_ops.convert_batch_matmul(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite BATCH_MATMUL → ONNX MatMul with optional transposes.

yobx.litert.ops.nn_ops.convert_conv2d(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite CONV_2D → ONNX Conv.

TFLite weight layout: (out_ch, kH, kW, in_ch) — convert to ONNX (out_ch, in_ch, kH, kW) via a Transpose. TFLite input/output layout: NHWC — wrapped with Transpose nodes to present as NCHW to the ONNX Conv node.

yobx.litert.ops.nn_ops.convert_depthwise_conv2d(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite DEPTHWISE_CONV_2D → ONNX Conv with group=in_channels.

TFLite weight layout: (1, kH, kW, out_ch) — for a depthwise conv with one filter per input channel, the number of groups equals the number of input channels. ONNX expects (out_ch, 1, kH, kW) for a grouped conv.

yobx.litert.ops.nn_ops.convert_fully_connected(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite FULLY_CONNECTED → ONNX MatMul (+ optional Add bias).

TFLite stores the weight matrix in (out_features, in_features) order so we transpose it to get the standard ONNX layout.

yobx.litert.ops.nn_ops.convert_max_pool2d(g: GraphBuilderExtendedProtocol, sts: Dict[str, Any], outputs: List[str], op: TFLiteOperator) → str[source]#

TFLite MAX_POOL_2D → ONNX MaxPool.