Skip to content

axes

torch_to_nnef.nemo_tract.axes

Dynamic axes helpers for NeMo Tract wrappers.

collapse_dynamic_axes_mapping

collapse_dynamic_axes_mapping(nemo_dynamic_axes: T.Dict[str, T.Dict[int, str]], input_names: T.Sequence[str]) -> T.Dict[str, T.Dict[int, str]]

Collapse mapping for all inputs keeping only the exposed inputs.