(tvm) 01/03: Preserve explicit multiplication in Relax shape syntax

[tq...@apache.org]

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 <tq...@users.noreply.github.com>
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