config

torch_to_nnef.nemo_tract.config

Configuration dataclasses for NeMo→NNEF export.

Defines lightweight, structured containers for both the CLI (NemoTractConfig and sub-configs) and the programmatic API (NemoExportConfig). Kept separate from runtime logic to minimise dependencies and improve clarity.

Lightweight enums that other modules need at import time (e.g. InspectFormat) live here to avoid circular imports.

CompressionConfig dataclass

CompressionConfig(compress_registry: str = DEFAULT_COMPRESSION_REGISTRY, compress_method: T.Optional[str] = None, dump_checked_io: bool = False)

Compression options for output artifacts.

InspectFormat

Bases: Enum

Output formats for inspection results.

InspectionConfig dataclass

InspectionConfig(inspect_signatures: bool = False, inspect_stages: T.Optional[T.List[str]] = None, inspect_format: InspectFormat = InspectFormat.HUMAN_RICH, inspect_output: T.Optional[Path] = None, inspect_diff: bool = False, shape_config: T.Optional[Path] = None, dump_shape_config: T.Optional[Path] = None, dry_run: bool = False)

Inspection and dry-run options.

LogConfig dataclass

LogConfig(verbose: bool = False)

Logging verbosity.

ModelSelectionConfig dataclass

ModelSelectionConfig(model_slug: str = '*', model_path: T.Optional[str] = None)

Model selection arguments.

NamingPrecisionConfig dataclass

NamingPrecisionConfig(naming_scheme: VariableNamingScheme = VariableNamingScheme.NATURAL_VERBOSE_CAMEL, data_type: str = 'float32')

Naming scheme and precision.

NemoExportConfig dataclass

NemoExportConfig(subnet: SubnetSelectionConfig = SubnetSelectionConfig(), compression: CompressionConfig = CompressionConfig(), naming_scheme: VariableNamingScheme = VariableNamingScheme.NATURAL_VERBOSE_CAMEL, data_type: str = 'float32', pretrained_name: str = '')

Configuration for the export_nemo_from_model programmatic API.

All fields have sensible defaults so callers only specify what they actually need to override.

Attributes:

Name	Type	Description
subnet	SubnetSelectionConfig	Subnet filtering and composition.
compression	CompressionConfig	Compression options for NNEF artifacts.
naming_scheme	VariableNamingScheme	NNEF variable naming scheme.
data_type	str	Export dtype policy: float32 | float16 | mixed.
pretrained_name	str	Human-readable label embedded in export metadata.

NemoTractConfig dataclass

NemoTractConfig(model: ModelSelectionConfig = ModelSelectionConfig(), output: OutputConfig = OutputConfig(), subnet: SubnetSelectionConfig = SubnetSelectionConfig(), sdpa: SdpaConfig = SdpaConfig(), tract: TractBinaryConfig = TractBinaryConfig(), naming: NamingPrecisionConfig = NamingPrecisionConfig(), compression: CompressionConfig = CompressionConfig(), inspect: InspectionConfig = InspectionConfig(), log: LogConfig = LogConfig())

Container for CLI arguments for NeMo→NNEF export.

Attributes:

Name	Type	Description
model	ModelSelectionConfig	Model selection arguments.
output	OutputConfig	Export destination arguments.
subnet	SubnetSelectionConfig	Subnet filtering arguments.
sdpa	SdpaConfig	SDPA-related export toggles.
tract	TractBinaryConfig	Tract binary path/version and IO-check.
naming	NamingPrecisionConfig	Naming scheme and precision.
compression	CompressionConfig	Compression options.
inspect	InspectionConfig	Inspection and dry-run options.
log	LogConfig	Logging verbosity.

OutputConfig dataclass

OutputConfig(export_dir: Path = Path(''))

Output destination arguments.

SdpaConfig dataclass

SdpaConfig(force_sdpa_pytorch: bool = False, tract_reify_sdpa: bool = False)

SDPA-related export toggles.

SubnetSelectionConfig dataclass

SubnetSelectionConfig(skip_preprocessor: bool = False, split_joint_decoder: bool = False, only_subnets: T.Optional[T.List[str]] = None)

Subnet filtering and composition.

TractBinaryConfig dataclass

TractBinaryConfig(tract_specific_version: T.Optional[str] = None, tract_specific_path: T.Optional[str] = None, tract_check_io_tolerance: T.Union[str, TractCheckTolerance] = TractCheckTolerance.APPROXIMATE.value)

Tract binary selection and IO-check policy.