vollo_torch.fx

vollo_torch.fx.prepare_shape(model: Module, *inputs: Tensor) → Tuple[GraphModule, Tensor | Tuple[Tensor, ...]]

Trace the model’s execution to annotate it with activation shapes. Necessary before converting a model to NNIR.

vollo_torch.fx.save(module: Module, input: Tensor, archive: str | Path)

Trace and save the Torch FX representation of a PyTorch model to an archive file.

This wraps torch.fx.GraphModule.to_folder to make it only use relative paths in the saved code, making it more portable.

Parameters:

  • module – PyTorch model to save.

  • input – Input tensor for the model, needed to trace the model with concrete shapes.

  • archive – Archive file to save model as. The arhive will be a valid .tar.gz file.

vollo_torch.fx.load(archive: str | Path) → Tuple[Module, Size]

Load a PyTorch model that has been saved to a folder with save().

Parameters:

archive – Archive file that a model has been saved as.

Returns:

A tuple (module, input_shape) where module is the loaded PyTorch model and input_shape is the shape of the input that the model was traced with.

vollo_torch.fx.is_vollo_tracing: ContextVar[bool] = <ContextVar name='is_vollo_tracing' default=False>

Describes whether the Vollo tracer is currently active. Make sure to use .get() to read the value.