prune_reshape_ops

Classes

PruneReshapeOps

PruneReshapeOps(enabled: bool = True)

The purpose of this transform is to prune RESHAPE operators.

1. Find any RESHAPE 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