Source code for pycif.plugins.modes.training.serialization.onnx_export

"""ONNX export helper.

Thin wrapper over :meth:`SurrogateBackend.export_onnx` so callers can use
the serialisation API without holding a reference to the backend object.
"""
from __future__ import annotations

from typing import Any, Tuple




def export_to_onnx(
    backend: Any,
    path: str,
    model: Any,
    input_shape: Tuple[int, ...],
) -> None:
    """Export *model* to ONNX using *backend*.

    Args:
        backend: A concrete
            :class:`~pycif.plugins.modes.training.backends.base.SurrogateBackend`.
        path: Destination file path, e.g. ``"surrogate.onnx"``.
        model: The trained framework-native model.
        input_shape: Shape of a single input tensor including the batch dim,
            e.g. ``(1, 64)``.
    """
    backend.export_onnx(path, model, input_shape)