Skip to content

Commit cc707ac

Browse files
authored
[CINN] Processing the null tensor output of IF op (#71915)
* [CINN]Improve the infer symbol of control flow op * support NullShapeData, 0D Tensor, TensorListShapeData * correct error messages * fix bug * fix null with tensorShapeData * fix bug
1 parent 9163673 commit cc707ac

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc

+5
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,11 @@ bool IfOp::InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context) {
445445
}
446446
infer_context->SetShapeOrDataForValue(
447447
result(rst_idx), symbol::ShapeOrDataDimExprs{result_list});
448+
} else if (false_shape_data.isa<symbol::NullShapeOrDataDimExpr>()) {
449+
infer_context->SetShapeOrDataForValue(result(rst_idx), true_shape_data);
450+
} else if (true_shape_data.isa<symbol::NullShapeOrDataDimExpr>()) {
451+
infer_context->SetShapeOrDataForValue(result(rst_idx),
452+
false_shape_data);
448453
} else {
449454
PADDLE_THROW(common::errors::Unimplemented(
450455
"IfOp::InferSymbolicShape: now only support "

0 commit comments

Comments
 (0)