onnx_diagnostic.torch_export_patches.onnx_export_serialization

onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_cache_serialization(patch_transformers: bool = False, patch_diffusers: bool = True, verbose: int = 0) → Dict[str, bool]

Registers many classes with onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization(). Returns information needed to undo the registration.

Parameters:

• patch_transformers – add serialization function for transformers package

• patch_diffusers – add serialization function for diffusers package

• verbosity – verbosity level

Returns:

information to unpatch

onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization(cls, f_flatten: Callable, f_unflatten: Callable, f_flatten_with_keys: Callable, f_check: Callable | None = None, verbose: int = 0) → bool

Registers a class. It can be undone with onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization().

Parameters:

• cls – class to register

• f_flatten – see torch.utils._pytree.register_pytree_node

• f_unflatten – see torch.utils._pytree.register_pytree_node

• f_flatten_with_keys – see torch.utils._pytree.register_pytree_node

• f_check – called to check the registration was successful

• verbose – verbosity

Returns:

registered or not

onnx_diagnostic.torch_export_patches.onnx_export_serialization.serialization_functions(patch_transformers: bool = False, patch_diffusers: bool = False, verbose: int = 0) → Dict[type, Callable[[int], bool]]

Returns the list of serialization functions.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0)

Undo all registrations.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization(cls: type, verbose: int = 0)

Undo the registration.