Skip to content

torch_to_nnef.remodeler.adapter

Classes:

Name Description
BoundaryAdapter

Boundary adapter applying tuple flattening and collapse at export time.
RenameOutputs

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

BoundaryAdapter 

BoundaryAdapter(module: Module, subnet_name: str, input_example: list, dynamic_axes: Optional[dict[str, dict[int, str]]], collapse_by_input: Optional[dict[str, set[str]]], binds_by_input: Optional[dict[str, str]] = None, renamed_map: Optional[dict[str, list[str]]] = None, outputs_keep: Optional[list[str]] = None, output_collapse_dims: 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.

RenameOutputs 

RenameOutputs(module: Module, rename_map: 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').