Skip to content

torch_to_nnef.remodeler.dyn_axes

Provider-agnostic helpers for dynamic-axes manipulation.

These utilities operate on the generic dyn mapping ({input_name: {axis_index: symbol}}) and assertion/extension strings. They are used by NeMo and can be reused by any provider.

Functions:

Name Description
apply_eval_symbols

Resize test_input tensors according to eval_symbols.

apply_symbol_renames_to_dyn

Apply symbol renames directly to a dynamic axes mapping.

drop_assertions_referencing_symbols

Drop assertions that name one of removed_symbols verbatim.

filter_assertions_present_in_dyn

Drop assertions that reference a symbol absent from every axis.

remove_eval_symbols_from_dyn

Remove pinned axes from dyn so the backend treats them as constant.

rewrite_and_filter_assertions

Rewrite assertions and drop those referencing removed symbols.

rewrite_assertions_with_renames

Rewrite assertion symbol names based on a rename mapping.

apply_eval_symbols

apply_eval_symbols(test_input: list, input_names: list[str], subnet_name: str, dyn: Dict[str, Dict[int, str]], eval_symbols: Dict[str, Dict[str, int]]) -> list

Resize test_input tensors according to eval_symbols.

apply_symbol_renames_to_dyn

apply_symbol_renames_to_dyn(dyn: Dict[str, Dict[int, str]], rename_map: Dict[str, List[str]]) -> T.Dict[str, T.Dict[int, str]]

Apply symbol renames directly to a dynamic axes mapping.

This is the lightweight alternative to BoundaryAdapter when only symbol renames are needed (no collapse, bind, or output filtering).

drop_assertions_referencing_symbols

drop_assertions_referencing_symbols(assertions: list[str], removed_symbols: Optional[Iterable[str]]) -> list[str]

Drop assertions that name one of removed_symbols verbatim.

Unlike :func:filter_assertions_present_in_dyn, this does not require every identifier in the assertion to be a known axis symbol: a hand-written or slug/derived-declared assertion may legitimately use a label (e.g. tract_assert tg: S==1) or reference something outside dyn entirely, and must not be discarded for that alone. It only drops an assertion when it references a symbol known to have just been pinned static (e.g. via eval_symbols), i.e. one that used to exist but provably no longer does.

Returns de-duplicated assertions.

filter_assertions_present_in_dyn

filter_assertions_present_in_dyn(assertions: list[str], dyn: Optional[dict[str, dict[int, str]]]) -> list[str]

Drop assertions that reference a symbol absent from every axis.

An assertion whose symbol(s) never appear in dyn (renamed away, collapsed, or bound) can never be evaluated against a real dimension: it is dead weight in the exported NNEF that tract itself flags as a "mislabeled symbol name" warning. Requires every identifier the assertion mentions to be a known axis symbol, so it only fits the single-symbol, machine-generated shape (e.g. tract_assert S >= 1) that :func:torch_to_nnef_nemo.dynaxes.build_dynamic_axes produces -- a free-form, user-declared assertion may use syntax this does not parse (a label, a function call) and should go through :func:drop_assertions_referencing_symbols instead.

Returns de-duplicated assertions.

remove_eval_symbols_from_dyn

remove_eval_symbols_from_dyn(input_names: list[str], subnet_name: str, dyn: Dict[str, Dict[int, str]], eval_symbols: Dict[str, Dict[str, int]]) -> None

Remove pinned axes from dyn so the backend treats them as constant.

Must be called after the BoundaryAdapter is built, because the adapter needs the symbols to resolve bindings.

rewrite_and_filter_assertions

rewrite_and_filter_assertions(assertions: list[str], rename_map: Optional[dict[str, list[str]]], dyn: Optional[dict[str, dict[int, str]]]) -> list[str]

Rewrite assertions and drop those referencing removed symbols.

  • Applies symbol renames so source symbols map to their target alias.
  • Computes the set of present symbols from the current dynamic axes and discards any assertion that mentions a symbol not present after rewriting.
  • Returns de-duplicated assertions.

rewrite_assertions_with_renames

rewrite_assertions_with_renames(assertions: list[str], rename_map: Optional[dict[str, list[str]]]) -> list[str]

Rewrite assertion symbol names based on a rename mapping.

Parameters:

Name Type Description Default

assertions

list[str]

List of assertion strings, e.g. "tract_assert U = BATCH".

required

rename_map

Optional[dict[str, list[str]]]

Mapping of target symbol to list of source symbols that should be rewritten to the target. Comparison is case-insensitive; rewritten symbols are emitted uppercased.

required

Returns:

Type Description
list[str]

A list of assertions with symbols rewritten according to

list[str]

the provided mapping. Unknown tokens are left unchanged.