torch_to_nnef.remodeler

Provider-agnostic boundary remodeler scaffold.

This module defines small, typed building blocks to describe IO signatures and boundary-only transforms (collapse, bind, and backend-facing symbol renames), plus helpers to load/save a strict nested config.

Notes: - The concrete YAML/JSON schema is identical to the NeMo remodeler and is parsed by the shared AxisSymbolRegistry loader used in the NeMo path. - Providers are expected to discover per-subnet signatures, and to apply a remodel plan by wrapping inner modules with an adapter that enforces the external boundary while preserving the internal contract.

BoundaryAdapter

BoundaryAdapter(module: torch.nn.Module, subnet_name: str, input_example: list, dynamic_axes: T.Optional[dict[str, dict[int, str]]], collapse_by_input: T.Optional[dict[str, set[str]]], binds_by_input: T.Optional[dict[str, str]] = None, renamed_map: T.Optional[dict[str, list[str]]] = None, outputs_keep: T.Optional[list[str]] = None, output_collapse_dims: T.Optional[dict[str, list[int]]] = None, *, apply_symbol_renames: bool = True)

Bases: Module

Boundary adapter applying tuple flattening and collapse at export time.

Both inputs AND outputs are flattened to flat tensor IO using :func:build_new_names_and_elements, so that the external interface (names, examples, dynamic axes) is always in terms of flat tensors.

When outputs_keep is set the adapter runs one forward pass at construction time to discover the output structure, then validates outputs_keep against the flattened output names and filters accordingly.

IODescriptor `dataclass`

IODescriptor(name: str, shape: list[T.Union[int, str]], dtype: T.Optional[str] = None, notes: T.Optional[list[str]] = None)

Description of a single input or output.

Provider

Bases: Protocol

Provider SPI to discover signatures and apply remodel plans.

Implementations should be small adapters around an existing provider (e.g., NeMo, plain PyTorch) that can: - discover raw and post-processed signatures for inspection - apply the boundary remodel plan to return wrapped modules ready for export

RemodelPlan `dataclass`

RemodelPlan(registry: 'AxisSymbolRegistry')

A remodel plan built from a validated AxisSymbolRegistry.

Attributes: - registry: The validated, parsed axis registry (nested schema).

RenameOutputs

RenameOutputs(module: torch.nn.Module, rename_map: T.Dict[str, str])

Bases: Module

Wrapper that renames output tensor names for export-time only.

Leaves computation unchanged and preserves input names. Useful to avoid name collisions between inputs and outputs (e.g., both named 'length').

Stage

Bases: Enum

Logical stages for signature inspection.

order `property`

order: int

Stable sort order (RAW < COLLAPSED < BOUND < FINAL).

SubnetSignature `dataclass`

SubnetSignature(name: str, stage: Stage, inputs: list[IODescriptor], outputs: list[IODescriptor], symbol_axes: T.Optional[dict[str, dict[int, str]]] = None, applied_flags: T.List[str] = list())

Per-subnet signature snapshot at a given stage.

plan_from_registry

plan_from_registry(registry: T.Any) -> RemodelPlan

Build a remodel plan from a validated registry (provider-specific).

save_config

save_config(path: T.Union[Path, str, None], registry: 'AxisSymbolRegistry', *, flow_seq: bool = True, stream: T.Optional[T.TextIO] = None) -> None

Save an AxisSymbolRegistry to YAML or JSON.

Parameters:

Name	Type	Description	Default
`path`	`Union[Path, str, None]`	Output file path (.yml/.yaml/.json) or None when using stream.	required
`registry`	`'AxisSymbolRegistry'`	Parsed registry to serialize.	required
`flow_seq`	`bool`	When YAML, render short lists in flow style.	`True`
`stream`	`Optional[TextIO]`	Optional text stream to write into (YAML or JSON). When provided, `path` can be None and the function infers YAML vs JSON based on the filename if available; defaults to YAML behavior.	`None`

torch_to_nnef.remodeler