[erhsh]

Description

Fixes #20928.

After the standalone TIRx extraction (#20803, #20831), many Python module paths moved from the parent tvm.* namespace into the tvm_tirx.* package, but documentation still referenced the removed paths, so executable doc examples fail with ModuleNotFoundError/AttributeError. The Python tutorials (e.g. docs/deep_dive/relax/tutorials/relax_creation.py) and the test suite were already migrated; this PR brings the .rst documentation and tutorial comments in line with the spellings they use.

Replacements

Stale spelling	Replacement
from tvm.script import relax as R	from tvm.relax import script as R
from tvm.script import s_tir as Ts	from tvm.s_tir import script as Ts
from tvm.script import ir as I	from tvm_tirx.script import ir as I
from tvm.script import tirx as T / Tx	from tvm_tirx.script import tirx as T / Tx
@tvm.script.ir_module	@I.ir_module (with the ir import added)
tvm.script.parse / from_source / register_dialect / parser	same names under tvm_tirx.script
tvm.tirx.compile / .layout / transform / Var / Add	tvm_tirx.tirx.*
tvm.ir.transform.* (PassContext, Sequential, instruments)	tvm_tirx.ir.transform.*
tvm.ir.{Call,Range,Stmt,register_op_attr,assert_structural_equal}	same under tvm_tirx.ir
from tvm.ir.type import ...	from tvm_tirx.ir.type import ...
tvm.target.Target / tvm.target.tag.*	tvm_tirx.target.*
tvm.IRModule	tvm_tirx.IRModule
tvm.backend.* / tvm.codegen	tvm_tirx.backend / tvm_tirx.codegen
tvm.runtime.{tensor,Tensor,empty,Device,load_module}	tvm_tirx.runtime.* (parent tvm.runtime only exports benchmarking/vm/disco)
device ctors tvm.cuda()/cpu()/metal()/vulkan()	tvm_tirx.* (incl. from_source.rst device checks)
import tvm, tvm.tirx (install verify command)	import tvm, tvm_tirx.tirx

Also:

• adds import tvm_tirx to example headers that now reference the package (tile_ops execution blocks, the CUDA tutorials, relax_vm.rst)
• fixes the stale tvm.base._LOADED_LIBS['tvm_runtime'] key in from_source.rst to tvm_runtime_extra (the key defined in python/tvm/base.py)
• fixes the docs/arch/relax_vm.rst example, which used R without importing it
• docs/reference/api/python/index.rst: the toctree caption is fixed to tvm_tirx.script (the autosummary pages it links to already document tvm_tirx.script)

Deliberately left unchanged

• TVMScript printer provenance comments (# from tvm.script import ...) inside printed output: they are non-executing comments and test expectations rely on them.
• The two historical "separate tvm.tir object model" mentions in docs/tirx/overview.rst and docs/tirx/arch/lowering_pipeline.rst: they contrast TIRx with the legacy TVM object model and are not live API references.
• docs/arch/index.rst prose that describes source folders (tvm/node, tvm/ir) — those sections describe the pre-extraction tree structure and deserve a dedicated rewrite; tracked in Docs reference nonexistent tvm.script / tvm.tirx paths after standalone TIRx extraction #20928.

Checklist

• PR title starts with [DOCS] per the commit tag convention
• ruff check / ruff format / whitespace hooks pass on all changed files (the check-asf-header hook raises a GBK UnicodeDecodeError on this Windows machine on pre-existing UTF-8 content and fails identically on unmodified files — a local locale artifact, not introduced by this change)

[erhsh]

Thanks for the review — all four findings are valid and are addressed in 43bfd6d:

1. [P2] missing tvm_tirx import in the standalone script — confirmed: the tutorial's own module imports tvm_tirx (line 68), but the run_mlp.py example embedded in comments did not. Added import tvm_tirx to its header. Also added it to the two runnable examples in docs/arch/index.rst that had the same gap.

2. [P2] wrong tmem_datapath_layout import path — confirmed: the function is defined in tvm_tirx/backend/cuda/layout.py and is not re-exported from tvm_tirx.tirx.layout; the first pass rewrote the prefix mechanically. Fixed to from tvm_tirx.backend.cuda.layout import tmem_datapath_layout, matching docs/tirx/tile_ops/tcgen05_ld_st.rst.

3. [P2] serialization example missing the relax import — confirmed: the first pass replaced from tvm import relax with the script-alias imports in the leading example, breaking relax.VirtualMachine further down. The example now imports VirtualMachine explicitly (same style as the neighboring "Runtime Execution" block); also added import tvm_tirx to the later usage block that needed it.

4. [P3] Tensor annotation lowercased — confirmed root cause: the bulk replacement ran through PowerShell -replace, which is case-insensitive, so the tvm.runtime.tensor rule consumed the capitalized Tensor annotations before the case-specific rule could apply. Restored tvm_tirx.runtime.Tensor in both places and verified no other annotation-style victims remain (git grep -E ":\s*tvm_tirx\.runtime\.tensor\b" -- docs is empty).