vollo_torch.fx.nnir

Translate Torch models to Vollo compiler NNIR models

vollo_torch.fx.nnir.to_nnir(module: GraphModule, inputs_precisions: NumberFormat | Sequence[NumberFormat] = NumberFormat.BF16, outputs_precisions: NumberFormat | Sequence[NumberFormat] = NumberFormat.BF16)

Note: Expects module to have gone through torch.fx.passes.shape_prop.ShapeProp().

Does not support inplace or _out variants of torch modules and functions.

Parameters:

• inputs_precisions – Specifies the precision of each model input. Supplying a non-Sequence value gives that precision to all inputs.

• outputs_precisions – Specifies the precision of each model output. Supplying a non-Sequence value gives that precision to all outputs.