File tree Expand file tree Collapse file tree 3 files changed +12
-4
lines changed
onnx_diagnostic/torch_export_patches Expand file tree Collapse file tree 3 files changed +12
-4
lines changed Original file line number Diff line number Diff line change @@ -114,7 +114,10 @@ Here is the list of supported caches:
114114
115115 import onnx_diagnostic.torch_export_patches.onnx_export_serialization as p
116116
117- print("\n ".join(sorted(t.__name__ for t in p.serialization_functions())))
117+ print(
118+ "\n ".join(sorted(t.__name__ for t in p.serialization_functions(
119+ patch_transformers=True, patch_diffusers=True)))
120+ )
118121
119122.. _l-control-flow-rewriting :
120123
Original file line number Diff line number Diff line change @@ -14,7 +14,10 @@ The following code shows the list of serialized classes in transformers.
1414
1515 import onnx_diagnostic.torch_export_patches.onnx_export_serialization as p
1616
17- print('\n '.join(sorted(t.__name__ for t in p.serialization_functions())))
17+ print(
18+ '\n '.join(sorted(t.__name__ for t in p.serialization_functions(
19+ patch_transformers=True, patch_diffusers=True
20+ ))))
1821
1922Patched Classes
2023===============
Original file line number Diff line number Diff line change @@ -28,7 +28,8 @@ def register_class_serialization(
2828) -> bool :
2929 """
3030 Registers a class.
31- It can be undone with :func:`unregister_class_serialization`.
31+ It can be undone with
32+ :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.unregister_class_serialization`.
3233
3334 :param cls: class to register
3435 :param f_flatten: see ``torch.utils._pytree.register_pytree_node``
@@ -77,7 +78,8 @@ def register_cache_serialization(
7778 patch_transformers : bool = False , patch_diffusers : bool = True , verbose : int = 0
7879) -> Dict [str , bool ]:
7980 """
80- Registers many classes with :func:`register_class_serialization`.
81+ Registers many classes with
82+ :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_class_serialization`.
8183 Returns information needed to undo the registration.
8284
8385 :param patch_transformers: add serialization function for
You can’t perform that action at this time.
0 commit comments