def save_model(
        model: nn.Module,
        model_configs: Dict[str, Any],
        checkpoint_dir: str,
        optimizer: Optional[torch.optim.Optimizer] = None,
        epoch: Optional[int] = None,
    ) -> None:
    """Save a PyTorch model checkpoint.

    The model weights are saved as `weights.pt` and the model
    configuration is saved as `configs.json` inside the provided
    checkpoint directory.

    Parameters
    ----------
    model : nn.Module
        The PyTorch model to save.

    model_configs : Dict[str, Any]
        Dictionary containing the parameters required to recreate the model.

    checkpoint_dir : str
        Directory where the checkpoint files will be stored.

    optimizer : Optional[torch.optim.Optimizer], default=None
        Optimizer whose state will be saved alongside the model. When
        set to `None`, the optimizer state will not be saved.

    epoch : Optional[int], default=None
        Current training epoch. If provided, it will be saved alongside the model.

    Returns
    -------
    None

    Examples
    --------
    Save a trained model and its configuration:

    >>> import torch
    >>> model = torch.nn.Linear(10, 2)
    >>> configs = {
    ...     "input_features": 10,
    ...     "output_features": 2,
    ... }
    >>> save_model(
    ...     model=model,
    ...     model_configs=configs,
    ...     checkpoint_dir="checkpoints/model_v1",
    ... )

    The resulting directory contains:

    .. code-block:: text

        checkpoints/model_v1/
        ├── weights.pt
        └── configs.json
    """

    model_path = os.path.join(checkpoint_dir, "weights.pt")
    checkpoint_path = os.path.join(checkpoint_dir, "configs.json")

    dict_to_save = {
        "model_state_dict": model.state_dict(),
    }
    if optimizer is not None:
        dict_to_save["optimizer_state_dict"] = optimizer.state_dict()
    if epoch is not None:
        dict_to_save["epoch"] = epoch

    torch.save(dict_to_save, model_path)
    save_file(model_configs, checkpoint_path)

def load_model(
        model_class: nn.Module,
        checkpoint_dir: str,
        map_location: str = "cpu",
        configs_to_override: Optional[Dict[str, Any]] = None,
    ) -> Tuple[nn.Module, Optional[Dict[str, Any]], Optional[int]]:
    """Load a PyTorch model checkpoint.

    The model class must accept the parameters stored in
    `configs.json` as keyword arguments.

    Parameters
    ----------
    model_class : nn.Module
        Model class used to recreate the architecture.

    checkpoint_dir : str
        Directory containing `weights.pt` and `configs.json`.

    map_location : str, default="cpu"
        Device where model weights are loaded.

    configs_to_override : Optional[Dict[str, Any]], default=None
        Configuration values that replace those saved in `configs.json` before
        the model is initialized. Only checkpoint tensors whose names and
        shapes match the resulting model are loaded.

    Returns
    -------
    Tuple[nn.Module, Optional[Dict[str, Any]], Optional[int]]
        Reconstructed model with loaded weights, optimizer state, and epoch.
        Missing optimizer state or epoch values are returned as `None`.

    Examples
    --------
    Load a saved model:

    >>> import torch
    >>> class MyModel(torch.nn.Module):
    ...     def __init__(self, input_features, output_features):
    ...         super().__init__()
    ...         self.layer = torch.nn.Linear(
    ...             input_features,
    ...             output_features,
    ...         )
    ...
    >>> model, optimizer_state, epoch = load_model(
    ...     model_class=MyModel,
    ...     checkpoint_dir="checkpoints/model_v1",
    ... )
    """

    if not os.path.exists(checkpoint_dir):
        logger.error(f"Checkpoint directory '{checkpoint_dir}' does not exist.", FileNotFoundError)

    model_configs_path = os.path.join(checkpoint_dir, "configs.json")
    if not os.path.exists(model_configs_path):
        logger.error(f"Model configuration file '{model_configs_path}' does not exist.", FileNotFoundError)

    model_path = os.path.join(checkpoint_dir, "weights.pt")
    if not os.path.exists(model_path):
        logger.error(f"Model weights file '{model_path}' does not exist.", FileNotFoundError)

    model_configs = load_file(model_configs_path)
    if configs_to_override is not None:
        model_configs.update(configs_to_override)

    model_init_signature = inspect.signature(model_class)
    accepts_arbitrary_kwargs = any(
        parameter.kind == inspect.Parameter.VAR_KEYWORD
        for parameter in model_init_signature.parameters.values()
    )

    if not accepts_arbitrary_kwargs:
        accepted_kwargs = {
            name
            for name, parameter in model_init_signature.parameters.items()
            if parameter.kind in {
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
                inspect.Parameter.KEYWORD_ONLY,
            }
        }
        ignored_configs = {
            name: value
            for name, value in model_configs.items()
            if name not in accepted_kwargs
        }

        if ignored_configs:
            logger.warning(
                f"Ignoring model config kwargs not accepted by "
                f"{model_class.__name__}: {ignored_configs}"
            )
            model_configs = {
                name: value
                for name, value in model_configs.items()
                if name in accepted_kwargs
            }

    model: nn.Module = model_class(**model_configs)
    checkpoint = torch.load(model_path, map_location=map_location)
    checkpoint_state_dict = checkpoint.get("model_state_dict", checkpoint)
    optimizer_state_dict = checkpoint.get("optimizer_state_dict")
    epoch = checkpoint.get("epoch")
    model_state_dict = model.state_dict()

    compatible_state_dict = {
        name: tensor
        for name, tensor in checkpoint_state_dict.items()
        if name in model_state_dict and tensor.shape == model_state_dict[name].shape
    }


    ### Print model table of layer status ###

    rows = []
    for name in sorted(set(checkpoint_state_dict) | set(model_state_dict)):
        checkpoint_tensor = checkpoint_state_dict.get(name)
        model_tensor = model_state_dict.get(name)

        if name in compatible_state_dict:
            status = "Loaded"
        elif checkpoint_tensor is None:
            status = "Initialized"
        elif model_tensor is None:
            status = "Ignored"
        else:
            status = "Incompatible"

        rows.append({
            "Layer": name,
            "Checkpoint shape": None if checkpoint_tensor is None else tuple(checkpoint_tensor.shape),
            "Model shape": None if model_tensor is None else tuple(model_tensor.shape),
            "Status": status,
        })

    print_table(pd.DataFrame(rows).set_index("Layer"), tablefmt="simple")
    model.load_state_dict(compatible_state_dict, strict=False)

    return model, optimizer_state_dict, epoch

def export_to_onnx(
        model: nn.Module,
        dummy_input: List[torch.Tensor],
        input_names: List[str],
        output_names: List[str],
        export_path: str,
        opset_version: int = 18,
        external_data: bool = False,
    ) -> None:
    """Export a PyTorch model to ONNX format and verify the conversion.

    After exporting, the generated ONNX model is executed using ONNX Runtime
    and its outputs are compared against the original PyTorch model outputs.

    Parameters
    ----------
    model : nn.Module
        PyTorch model to export.

    dummy_input : List[torch.Tensor]
        List of dummy input tensors to be passed to the model.

    input_names : List[str]
        Names assigned to ONNX input tensors. Must match the order of `dummy_input`.

    output_names : List[str]
        Names assigned to ONNX output tensors.

    export_path : str
        Destination path for the ONNX file.

    opset_version : int, default=18
        ONNX opset version used during export.

    external_data : bool, default=False
        Store model weights externally. Useful for models larger than 2GB.

    Returns
    -------
    None

    Examples
    --------
    Export a simple neural network:

    >>> import torch
    >>> model = torch.nn.Linear(10, 2)
    >>> export_to_onnx(
    ...     model=model,
    ...     dummy_input=[torch.randn(1, 10)],
    ...     input_names=["features"],
    ...     output_names=["prediction"],
    ...     export_path="model.onnx",
    ... )

    The generated ONNX model can then be executed using
    :func:`run_onnx`.
    """

    if len(dummy_input) != len(input_names):
        logger.error("The number of dummy inputs must match the number of input names.", ValueError)

    model.eval()

    torch.onnx.export(
        model, 
        tuple(dummy_input), 
        export_path, 
        opset_version=opset_version,
        input_names=input_names,
        output_names=output_names,
        external_data=external_data,
    )


    ### Check nn.Module and ONNX outputs ###

    with torch.no_grad():
        # Get PyTorch model output
        torch_output = model(*dummy_input)

        # Convert to tuple
        if isinstance(torch_output, torch.Tensor):
            torch_output = (torch_output,)

        # Remove None outputs since ONNX does not support None outputs
        # and removes them at runtime
        torch_output = tuple(output for output in torch_output if output is not None)

        # Load ONNX model and get output
        ort_session = ort.InferenceSession(export_path)
        ort_inputs = {name: input_tensor.numpy() for name, input_tensor in zip(input_names, dummy_input)}
        onnx_output = ort_session.run(None, ort_inputs)

        for torch_out, onnx_out in zip(torch_output, onnx_output):
            np.testing.assert_allclose(
                torch_out.numpy(), 
                onnx_out, 
                rtol=1e-03, 
                atol=1e-05,
            )

def run_onnx(
        onnx_model_path: str,
        input_data: Dict[str, np.ndarray],
    ) -> List[np.ndarray]:
    """Run an ONNX model using ONNX Runtime.

    Parameters
    ----------
    onnx_model_path : str
        Path to the ONNX model file.

    input_data : Dict[str, np.ndarray]
        Dictionary mapping ONNX input names to NumPy input arrays.

    Returns
    -------
    List[np.ndarray]
        Model output tensors.

    Examples
    --------
    Run inference on an exported ONNX model:

    >>> import numpy as np
    >>> outputs = run_onnx(
    ...     onnx_model_path="model.onnx",
    ...     input_data={
    ...         "features": np.random.randn(1, 10).astype(np.float32)
    ...     },
    ... )
    >>> prediction = outputs[0]
    """

    if not os.path.exists(onnx_model_path):
        logger.error(f"ONNX model file '{onnx_model_path}' does not exist.", FileNotFoundError)

    ort_session = ort.InferenceSession(onnx_model_path)
    outputs = ort_session.run(None, input_data)

    return outputs

def print_model_summary(model: nn.Module) -> None:
    """Print a summary of a PyTorch model.

    The summary includes the model architecture and total/trainable
    parameter counts.

    Parameters
    ----------
    model : nn.Module
        PyTorch model to summarize.

    Examples
    --------
    Print the summary of a model:

    >>> import torch
    >>> model = torch.nn.Sequential(
    ...     torch.nn.Linear(10, 5),
    ...     torch.nn.ReLU(),
    ...     torch.nn.Linear(5, 1),
    ... )
    >>> print_model_summary(model=model)
    """

    header_layer = "Layer (type)"
    header_params = "Param #"
    header_trainable = "Trainable"

    rows: List[Tuple[str, int, str]] = []
    total_params = 0
    trainable_params = 0

    for name, module in model.named_modules():
        # Skip the top-level container itself
        if name == "":
            continue

        # Only leaf modules (no children) to avoid double-counting
        if list(module.children()):
            continue

        layer_type = type(module).__name__
        label = f"{name} ({layer_type})"

        layer_total = sum(p.numel() for p in module.parameters())
        layer_trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)

        total_params += layer_total
        trainable_params += layer_trainable

        trainable_flag = "Yes" if layer_trainable > 0 else "No"
        rows.append((label, layer_total, trainable_flag))

    summary_labels = ["Total params", "Trainable params", "Non-trainable params"]
    summary_values = [total_params, trainable_params, total_params - trainable_params]

    # Size each column to fit its widest content so long names/values never wrap
    col_layer = max(len(header_layer), *(len(label) for label, _, _ in rows), *(len(label) for label in summary_labels))
    col_params = max(len(header_params), *(len(f"{value:,}") for _, value, _ in rows), *(len(f"{value:,}") for value in summary_values))
    col_trainable = max(len(header_trainable), *(len(flag) for _, _, flag in rows))

    header = f"{header_layer:<{col_layer}} {header_params:>{col_params}} {header_trainable:>{col_trainable}}"
    separator = "-" * len(header)

    print(separator)
    print(header)
    print("=" * len(header))

    for label, layer_total, trainable_flag in rows:
        print(f"{label:<{col_layer}} {layer_total:>{col_params},} {trainable_flag:>{col_trainable}}")

    print("=" * len(header))
    for label, value in zip(summary_labels, summary_values):
        print(f"{label:<{col_layer}} {value:>{col_params},}")
    print(separator)