(tvm) 08/33: [REFACTOR][Script] Emit imports from each dialect owning package

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit 72ac640cf4748b57e2cfa621e981031de918d63c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:37:32 2026 +0000

    [REFACTOR][Script] Emit imports from each dialect owning package
---
 tirx/src/script/printer/printer.cc | 11 +++++++----
 tirx/src/tirx/ir/index_map.cc      |  2 +-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/tirx/src/script/printer/printer.cc b/tirx/src/script/printer/printer.cc
index 271260ab8e..a8ffd598e2 100644
--- a/tirx/src/script/printer/printer.cc
+++ b/tirx/src/script/printer/printer.cc
@@ -360,9 +360,9 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
             if (auto id = child.as<IdDoc>()) used.insert(id.value()->name);
           },
           &visited);
-      std::string loader = "tvm";
+      std::string loader = "tvm_tirx";
       for (size_t suffix = 1; used.count(loader); ++suffix)
-        loader = "tvm_" + std::to_string(suffix);
+        loader = "tvm_tirx_" + std::to_string(suffix);
       ffi::String json = ffi::json::Stringify(
           ffi::ToJSONGraph(metadata, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
       block.value()->stmts.insert(block.value()->stmts.begin(),
@@ -373,7 +373,7 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
                                                 ->Call({LiteralDoc::Str(json, std::nullopt)}),
                                             std::nullopt));
       header.push_back(
-          ffi::String(loader == "tvm" ? "import tvm\n" : "import tvm as " + loader + "\n"));
+          ffi::String(loader == "tvm_tirx" ? "import tvm_tirx\n" : "import tvm_tirx as " + loader + "\n"));
     } else {
       block.value()->stmts.push_back(
           CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
@@ -411,7 +411,10 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
       canonical.resize(canonical.size() - std::string(".prefix").size());
       if (!namespaces.count(canonical)) continue;
       ffi::String alias = config->GetExtraConfig<ffi::String>(key, fallback);
-      ffi::String import = "from tvm_tirx.script import " + canonical + " as " + std::string(alias);
+      ffi::String import =
+          canonical == "s_tir" || canonical == "relax"
+              ? "from tvm." + canonical + " import script as " + std::string(alias)
+              : "from tvm_tirx.script import " + canonical + " as " + std::string(alias);
       if (config->GetExtraConfig<bool>("ir.comment_imports", false) && !config->show_meta) {
         header.push_back(CommentDoc(import));
       } else {
diff --git a/tirx/src/tirx/ir/index_map.cc b/tirx/src/tirx/ir/index_map.cc
index 547e385035..28fb38565c 100644
--- a/tirx/src/tirx/ir/index_map.cc
+++ b/tirx/src/tirx/ir/index_map.cc
@@ -473,7 +473,7 @@ ffi::String IndexMapNode::ToPythonString(
   std::string inverse_lambda_expr =
       IndexMap2PythonLambdaExpr(inverse->initial_indices, inverse->final_indices);
   std::ostringstream oss;
-  oss << "tvm.tirx.IndexMap.from_func(" << lambda_expr
+  oss << "tvm_tirx.tirx.IndexMap.from_func(" << lambda_expr
       << ", inverse_index_map=" << inverse_lambda_expr << ")";
   return ffi::String(oss.str());
 }