not shrink rank of shape when unroll tensor iterator#36381
not shrink rank of shape when unroll tensor iterator#36381michal-miotk wants to merge 4 commits into
Conversation
github-actions
Bot
added
the
category: transformations
CuriousPanCake
left a comment
CuriousPanCake
left a comment
There was a problem hiding this comment.
Could you please provide a better title/description of the PR? Also, it would be good to see a test and an example of what's happening
michal-miotk
changed the title
CuriousPanCake
left a comment
CuriousPanCake
left a comment
There was a problem hiding this comment.
LGTM
@praasz please take a look
There was a problem hiding this comment.
Pull request overview
This PR fixes a shape-rank regression in the UnrollTensorIterator transformation when replacing the Loop body’s current_iteration parameter with a Constant, ensuring the replacement preserves the original parameter rank (avoiding accidental scalarization that can break shape-sensitive ops like Gather).
Changes:
- Update
UnrollTensorIteratorto create the replacementcurrent_iterationConstantusing the original parameter shape (instead of always using scalar{}). - Add a regression test that validates the replacement
Constantrank and resultingGatheroutput rank remain unchanged after unrolling.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/common/transformations/tests/control_flow/unroll_tensor_iterator_test.cpp |
Adds a regression test covering shape-rank preservation for current_iteration replacement and Gather output after unrolling. |
src/common/transformations/src/transformations/control_flow/unroll_tensor_iterator.cpp |
Adjusts current_iteration replacement Constant construction to use the parameter’s shape (rank-preserving behavior). |
| auto cur_iter_const = std::make_shared<v0::Constant>(param_to_delete->get_element_type(), | ||
| param_to_delete->get_shape(), | ||
| idx); |
There was a problem hiding this comment.
according to doc: Trip count: A scalar or 1D tensor with 1 element of int64 or int32 type specifying maximum number of iterations. source: https://docs.openvino.ai/2026/documentation/openvino-ir-format/operation-sets/operation-specs/infrastructure/loop-5.html I don't think that above code is needed.
There was a problem hiding this comment.
If this is Parameter it can be also defined as PartialShape::dynamic() which mean it can be 1D, but get_shape will not work.
@CuriousPanCake
Or the dynamic shape is not possible here?
Details:
gather
Tickets: