Skip to content

torch_to_nnef.torch_graph.dynamic_axes

Conservative per-tensor dynamic-axis tracking for symbolic-dim export.

Under dynamic_axes only some tensor axes are symbolic; the rest keep their traced constant size. Without knowing which is which, t2n forces every shape-derived value dynamic (see _if_dyn_shape_may_remove_resolved_dim), so a split whose sizes come from x.shape[-1] on a static axis fails to lower. This pass computes, per tensor, the set of axis indices that are dynamic, so aten::size on a provably-static axis can still fold.

Design (why it is sound):

  • The returned set OVER-approximates the truly-dynamic axes. An axis is reported static (absent) only when provably so, so a truly symbolic dimension is never baked to a constant.
  • The DEFAULT for any op without an explicit rule is "all axes dynamic". In particular the elementwise/broadcasting rule (which assumes right-aligned axis identity) is applied ONLY to a whitelist of real elementwise ops; an axis-reordering op (transpose, permute, ...) that is not explicitly handled therefore falls back to all-dynamic rather than being mis-mapped.
  • Rules are STRUCTURAL: dynamic-ness is decided from op structure (reshape target literals, axis arguments, right-aligned broadcasting), never by comparing traced shape values -- a dynamic axis is often traced as size 1 (e.g. batch), indistinguishable from a static 1 by value.

Functions:

Name Description
compute_dynamic_axis_map

Map tensor.name -> set(dynamic axis indices) (over-approximation).

size_query_is_dynamic

Whether aten::size(input_node, axis) reads a dynamic axis.

compute_dynamic_axis_map

compute_dynamic_axis_map(ir_graph, dynamic_axes_by_name: Dict[str, Dict[int, str]]) -> T.Dict[str, T.Set[int]]

Map tensor.name -> set(dynamic axis indices) (over-approximation).

size_query_is_dynamic

size_query_is_dynamic(dynamic_axis_map: Dict[str, Set[int]], input_node, axis: int) -> bool

Whether aten::size(input_node, axis) reads a dynamic axis.

Conservative: if the tensor is unknown to the map, treat as dynamic.