prune_squeeze_ops

Classes

PruneSqueezeOps

PruneSqueezeOps(enabled: bool = True)

The purpose of this transform is to prune SQUEEZE operators.

1. Find any SQUEEZE operator where subsequent operators dont rely on input tensor.
2. If found, set output tensor buffer to input tensor buffer and drop the op.

Functions

apply

apply(model: ModelT, subgraph: SubGraphT, operations: list[Type[AotOperator]]) -> tuple[litert.ModelT, litert.SubGraphT, list[Type[AotOperator]]]

Apply the transform.

Parameters:

• model

  (ModelT) –

  The model to transform.

• subgraph

  (SubGraphT) –

  The subgraph to transform.

• operations

  (list[Type[AotOperator]]) –

  The list of operations in the subgraph.

Returns:

• tuple[ModelT, SubGraphT, list[Type[AotOperator]]] –

  tuple[litert.ModelT, litert.SubGraphT, list[Type["AotOperator"]]]: The transformed model, subgraph, and list of kept operations.

Functions