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 |
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 |
|---|---|---|---|
|
list[str]
|
List of assertion strings, e.g. "tract_assert U = BATCH". |
required |
|
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. |