onnx_diagnostic.torch_export_patches.patch_inputs

onnx_diagnostic.torch_export_patches.patch_inputs.convert_dynamic_axes_into_dynamic_shapes(model: Module, args: Tuple[Any, ...] | None = None, kwargs: Dict[str, Any] | None = None, dynamic_axes: Dict[str, Dict[int, str]] | None = None, prefix_mapping: Dict[str, str] | None = None, verbose: int = 0) → Tuple[Tuple[Any, ...], Dict[str, Any], Dict[str, Any]]

Converts the input from an export to something torch.export.export() can handle.

Parameters:

• model – model to convert (used to extract the signature)

• args – positional arguments

• kwargs – named arguments

• dynamic_axes – dynamic axes

• prefix_mapping – prefix mapping

• verbose – verbosity

Returns:

(args, kwargs, dynamic shapes)

onnx_diagnostic.torch_export_patches.patch_inputs.use_dyn_not_str(dynamic_shapes: Any) → Any

Some functions returns dynamic shapes as string. This functions replaces them with torch.export.Dim.DYNAMIC.