hyrax.model_exporters
=====================

.. py:module:: hyrax.model_exporters


Attributes
----------

.. autoapisummary::

   hyrax.model_exporters.logger


Functions
---------

.. autoapisummary::

   hyrax.model_exporters.export_to_onnx
   hyrax.model_exporters._export_pytorch_to_onnx
   hyrax.model_exporters._export_pytorch_to_onnx_v2


Module Contents
---------------

.. py:data:: logger

.. py:function:: export_to_onnx(model, sample, config, ctx)

   Dispatching function to convert a ML framework model into an ONNX model.

   :param model: The model that was just trained using the ML framework. i.e. PyTorch
   :type model: ML framework model
   :param sample: This sample is the result of running a batch of data through the data
                  loader and the model's `to_tensor` function. It is used to compare the
                  output of the ONNX model against the output of the PyTorch model.
   :type sample: Tensor
   :param config: The parsed config file as a nested dict
   :type config: dict
   :param ctx: A context dictionary containing info needed for the conversion to ONNX.
   :type ctx: dict


.. py:function:: _export_pytorch_to_onnx(model, sample, output_filepath, opset_version)

   Specific implementation to convert PyTorch model to ONNX format. This uses
   the older (torch<2.9) export capabilities. And only supports up the opset
   version 20.

   :param model: The PyTorch model to be converted to ONNX format.
   :type model: torch.nn.Module
   :param sample: A sample of input data to the model. This is used to trace the model
                  during the export process.
   :type sample: np.ndarray or list of np.ndarray
   :param output_filepath: The file path where the ONNX model will be saved.
   :type output_filepath: pathlib.Path
   :param opset_version: The ONNX opset version to use for the export.
   :type opset_version: int


.. py:function:: _export_pytorch_to_onnx_v2(model, sample, output_filepath, opset_version)

   Currently unused.
   Specific implementation to convert PyTorch model to ONNX format using
   torch Dynamo export capabilities.

   :param model: The PyTorch model to be converted to ONNX format.
   :type model: torch.nn.Module
   :param sample: A sample of input data to the model. This is used to trace the model
                  during the export process.
   :type sample: np.ndarray or list of np.ndarray
   :param output_filepath: The file path where the ONNX model will be saved.
   :type output_filepath: pathlib.Path
   :param opset_version: The ONNX opset version to use for the export.
   :type opset_version: int