Skip to content

[12/17] DLA rank conversion #2964

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mgohil-png wants to merge 7 commits into from
Draft

[12/17] DLA rank conversion #2964

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
8e54406
[DLA_transforms_12] [12/17] DLA rank conversion
mgohil-png Oct 15, 2026
6225a08
Declare deferred output ranks in rank-conversion fixtures
mgohil-png Oct 26, 2026
56cb037
Preserve runtime initializer defaults during deferred rank preparation
mgohil-png Oct 27, 2026
d8e9ba7
Preserve native layout contracts in ordered DLA rank stages
mgohil-png Oct 27, 2026
b12b4fa
Validate DLA rank-collapse contexts and prepare implicit Slice axes
mgohil-png Oct 27, 2026
19a831e
Preserve native index controls and public fixed layouts in rank conve…
mgohil-png Oct 27, 2026
00d2cb6
Preserve symbolic outputs and legacy rank controls
mgohil-png Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,8 +154,23 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False):
boundary_names = cache.graph_input_names | cache.init_names
append = []
plans = []
inputs_promotable = all(
cache.get_shape(value.name) is not None
for value in graph.input
if value.name not in cache.init_names
)
for value in graph.output:
shape = cache.get_shape(value.name)
shape: list | None = cache.get_shape(value.name)
if (
shape is None
and defer_rank_conversion
and inputs_promotable
and value.type.tensor_type.HasField("shape")
):
shape = [
dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None
for dim in value.type.tensor_type.shape.dim
]
Comment thread
coderabbitai[bot] marked this conversation as resolved.
if shape is None or len(shape) >= 4 or value.name in boundary_names:
continue
producer = cache.get_producer(value.name)
Expand Down Expand Up @@ -194,7 +209,7 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False):
_require_adapter_opset(model)
for value, shape, existing, axes in plans:
axes_name = None
if len(shape) > 1 or shape == [1]:
if len(shape) > 1 or shape == [1] or any(not isinstance(dim, int) for dim in shape):
axes_name = add_initializer(
graph,
cache.unique_name(value.name + "_squeeze_axes"),
Expand Down
Loading
Loading