documentation

Author: xadupre

_doc/patches.rst

@@ -104,7 +104,7 @@ and triggered by ``with torch_export_patches(patch_transformers=True)``.
 This function does one class, 
 :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_cache_serialization`
 does all known classes.
-It can be undone with :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister`
+It can be undone with :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization`
 or :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_cache_serialization`.
 Here is the list of supported caches:
 

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -28,7 +28,8 @@ def register_class_serialization(
 ) -> bool:
     """
     Registers a class.
-    It can be undone with :func:`unregister`.
+    It can be undone with
+    :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization`.
 
     :param cls: class to register
     :param f_flatten: see ``torch.utils._pytree.register_pytree_node``
@@ -77,7 +78,8 @@ def register_cache_serialization(
     patch_transformers: bool = False, patch_diffusers: bool = True, verbose: int = 0
 ) -> Dict[str, bool]:
     """
-    Registers many classes with :func:`register_class_serialization`.
+    Registers many classes with
+    :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization`.
     Returns information needed to undo the registration.
 
     :param patch_transformers: add serialization function for