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 |
size_query_is_dynamic |
Whether |