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][source][source]

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[source][source]

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]][source][source]

Returns the list of serialization functions.

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

Undo all registrations.

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

Undo the registration.