Introduce ReplaceWhileOperandShape api to HloInstruction. This function provides an interface to change the shape an tuple operand in while loops. #67127
Introduce ReplaceWhileOperandShape api to HloInstruction. This function provides an interface to change the shape an tuple operand in while loops.
The default implementation (ShapeTransformer) propagates shape change through instructions in which the output shape is directly inferred from operands, namely, gte, tuple, and nested while loops.
As the result, if the changed operand of the while loop is only used by the mentioned instructions, the call to ReplaceWhileOperandShape guarantees validity of the hlo graph after shape replacement. Currently, the replace function simply bails and returns false if there are any users other than the mentioned instructions.
Introduce ReplaceWhileOperandShape api to HloInstruction. This function provides an interface to change the shape an tuple operand in while loops.
The default implementation (ShapeTransformer) propagates shape change through instructions in which the output shape is directly inferred from operands, namely, gte, tuple, and nested while loops.
As the result, if the changed operand of the while loop is only used by the mentioned instructions, the call to ReplaceWhileOperandShape guarantees validity of the hlo graph after shape replacement. Currently, the replace function simply bails and returns false if there are any users other than the mentioned instructions.