onnx_diagnostic.torch_export_patches.onnx_export_serialization

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_base_model_output(bo: BaseModelOutput) → Tuple[List[Any], Any]

Serializes a transformers.modeling_outputs.BaseModelOutput with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_dynamic_cache(dynamic_cache: DynamicCache) → Tuple[List[Any], Any]

Serializes a transformers.cache_utils.DynamicCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_encoder_decoder_cache(ec_cache: EncoderDecoderCache) → Tuple[List[Any], Any]

Serializes a transformers.cache_utils.EncoderDecoderCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_mamba_cache(mamba_cache: MambaCache) → Tuple[List[Any], Any]

Serializes a transformers.cache_utils.MambaCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_sliding_window_cache(cache: SlidingWindowCache) → Tuple[List[Any], Any]

Serializes a transformers.cache_utils.SlidingWindowCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_with_keys_base_model_output(bo: BaseModelOutput) → Tuple[List[Tuple[KeyEntry, Any]], Any]

Serializes a transformers.modeling_outputs.BaseModelOutput with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache) → Tuple[List[Tuple[KeyEntry, Any]], Any]

Serializes a transformers.cache_utils.DynamicCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_with_keys_encoder_decoder_cache(ec_cache: EncoderDecoderCache) → Tuple[List[Tuple[KeyEntry, Any]], Any]

Serializes a transformers.cache_utils.EncoderDecoderCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_with_keys_mamba_cache(cache: MambaCache) → Tuple[List[Tuple[KeyEntry, Any]], Any]

Serializes a transformers.cache_utils.MambaCache with python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.flatten_with_keys_sliding_window_cache(cache: SlidingWindowCache) → Tuple[List[Tuple[KeyEntry, Any]], Any]

Serializes a transformers.cache_utils.SlidingWindowCache with python objects.

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

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

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 unregister().

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(verbose: int = 0) → Dict[str, Callable | int]

Returns the list of serialization functions.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unflatten_base_model_output(values: List[Any], context: Any, output_type=None) → BaseModelOutput

Restores a transformers.modeling_outputs.BaseModelOutput from python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unflatten_dynamic_cache(values: List[Any], context: Any, output_type=None) → DynamicCache

Restores a transformers.cache_utils.DynamicCache from python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unflatten_encoder_decoder_cache(values: List[Any], context: Any, output_type=None) → EncoderDecoderCache

Restores a transformers.cache_utils.EncoderDecoderCache from python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unflatten_mamba_cache(values: List[Any], context: Any, output_type=None) → MambaCache

Restores a transformers.cache_utils.MambaCache from python objects.

onnx_diagnostic.torch_export_patches.onnx_export_serialization.unflatten_sliding_window_cache(values: List[Any], context: Any, output_type=None) → SlidingWindowCache

Restores a transformers.cache_utils.SlidingWindowCache from python objects.

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

Undo the registration.

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

Undo all registrations.