Skip to content

nexuml.core.export

nexuml.core.export

Export, reload, and selective checkpoint loading for trained pipelines.

LoadReport dataclass

Selective load result for package/state reuse.

TrainingReload dataclass

Reloaded package prepared for current-codebase training.

export_package

export_package(
    pipeline: CompiledPipeline,
    path: Path,
    metadata: dict[str, Any] | None = None,
    lightning_module: Any | None = None,
    trainer: Any | None = None,
    checkpoint_path: str | Path | None = None,
    include_modules: list[str] | None = None,
    source_metadata: dict[str, Any] | None = None,
) -> Path

Export a trained pipeline as a rich package-backed artifact directory.

Parameters:

Name Type Description Default
pipeline CompiledPipeline

Compiled pipeline to export.

required
path Path

Destination directory.

required
metadata dict[str, Any] | None

Optional provenance metadata merged into the artifact.

None
lightning_module Any | None

Optional Lightning module for checkpoint sidecars.

None
trainer Any | None

Optional Lightning trainer for training-state sidecars.

None
checkpoint_path str | Path | None

Optional source Lightning checkpoint to preserve.

None
include_modules list[str] | None

Optional glob patterns for additional source modules to intern (useful for dynamic imports invisible to torch.package).

None
source_metadata dict[str, Any] | None

Optional metadata describing the export source (e.g. CLI checkpoint path). Merged into metadata.

None

Returns:

Type Description
Path

Path to the created export directory.

load_weights

load_weights(
    pipeline: CompiledPipeline,
    source: str | Path,
    checkpoint: CheckpointLoadSpec | None = None,
    *,
    include: list[str] | None = None,
    exclude: list[str] | None = None,
    allow_missing: bool | None = None,
    allow_shape_mismatch: bool | None = None,
    freeze_loaded: bool | None = None,
) -> LoadReport

Selectively load weights into an already-compiled pipeline.

Returns:

Type Description
LoadReport

LoadReport summarising matched, missing, and excluded keys.

Raises:

Type Description
ValueError

On shape mismatch or unexpected missing keys when not allowed.

load_package

load_package(
    path: Path,
) -> tuple[
    CompiledPipeline, ResolvedConfig, dict[str, Any]
]

Reload an exported pipeline into the current codebase.

Returns:

Type Description
tuple[CompiledPipeline, ResolvedConfig, dict[str, Any]]

Tuple of (pipeline, resolved_config, metadata).

Raises:

Type Description
ValueError

If no scenario config is found in the artifact.

load_inference_package

load_inference_package(
    path: Path,
) -> tuple[
    CompiledPipeline, ResolvedConfig, dict[str, Any]
]

Load the packaged pipeline object directly from the torch.package artifact.

Returns:

Type Description
tuple[CompiledPipeline, ResolvedConfig, dict[str, Any]]

Tuple of (pipeline, resolved_config, metadata).

Raises:

Type Description
FileNotFoundError

If no package artifact exists at path.

load_package_for_training

load_package_for_training(
    path: Path,
    scenario: ScenarioSpec | None = None,
    checkpoint: CheckpointLoadSpec | None = None,
) -> TrainingReload

Reload a package into the current codebase for resume or fine-tuning.

Returns:

Type Description
TrainingReload

TrainingReload with pipeline, lightning module, scenario, and load report.

Raises:

Type Description
ValueError

If no scenario is provided and the artifact has no packaged config.

export_safetensors

export_safetensors(
    pipeline: CompiledPipeline,
    path: Path,
    include: list[str] | None = None,
    exclude: list[str] | None = None,
    metadata: dict[str, Any] | None = None,
) -> Path

Export pipeline weights as SafeTensors plus a JSON manifest.

Returns:

Type Description
Path

Path to the created .safetensors file.

export_onnx

export_onnx(
    pipeline: CompiledPipeline,
    path: Path,
    input_key: str | None = None,
    output_key: str = "reconstructed",
    opset_version: int = 18,
) -> Path

Export an inference-only ONNX graph for single-input pipelines.

Returns:

Type Description
Path

Path to the created .onnx file.

Raises:

Type Description
ImportError

If the onnxscript package is not installed.

infer

infer(
    pipeline: CompiledPipeline,
    x: TensorDict,
    y: TensorDict | None = None,
) -> TensorDict

Run inference on a pipeline (eval mode, no grad).

Returns:

Type Description
TensorDict

Output TensorDict from the pipeline.