This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-refinement in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 923cf65477fefcfee29aaf8fbe054ce48d28e00c Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 1 20:21:31 2026 +0000 Preserve explicit multiplication in Relax shape syntax --- src/relax/script/printer/tir.cc | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc index d2ae3d0177..9c15d81ee6 100644 --- a/src/relax/script/printer/tir.cc +++ b/src/relax/script/printer/tir.cc @@ -28,11 +28,15 @@ namespace script { namespace printer { namespace details { -ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim) { - ExprDoc translated = d->Translate(dim).value(); - if (auto mul = dim.as<prim::MulNode>()) { +namespace { + +ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim, const ExprDoc& translated) { + auto operation = translated.as<OperationDoc>(); + if (auto mul = dim.as<prim::MulNode>(); + mul && operation && operation.value()->kind == OperationDocNode::Kind::kMult) { ExprDoc compact = OperationDoc(OperationDocNode::Kind::kMult, - {RelaxShapeDim(d, mul->a), RelaxShapeDim(d, mul->b)}); + {RelaxShapeDim(d, mul->a, operation.value()->operands[0]), + RelaxShapeDim(d, mul->b, operation.value()->operands[1])}); d->RecordOrigin(compact, dim); return compact; } @@ -45,6 +49,12 @@ ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim) { return translated; } +} // namespace + +ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim) { + return RelaxShapeDim(d, dim, d->Translate(dim).value()); +} + } // namespace details } // namespace printer } // namespace script
