[TIRX] Expose analyzer-aware Python stmt/expr functors

python/tvm/arith/analyzer.py

@@ -108,7 +108,16 @@ class Analyzer:
     """
 
     def __init__(self):
-        _mod = _ffi_api.CreateAnalyzer()
+        self._init_from_module(_ffi_api.CreateAnalyzer())
+
+    @classmethod
+    def _from_module(cls, mod):
+        analyzer = cls.__new__(cls)
+        analyzer._init_from_module(mod)
+        return analyzer
+
+    def _init_from_module(self, mod):
+        _mod = mod
         self._const_int_bound = _mod("const_int_bound")
         self._const_int_bound_update = _mod("const_int_bound_update")
         self._const_int_bound_is_bound = _mod("const_int_bound_is_bound")

python/tvm/tirx/__init__.py

@@ -104,7 +104,12 @@
 from . import backend
 from . import stmt_functor
 
-from .functor import PyStmtExprVisitor, PyStmtExprMutator
+from .functor import (
+    PyStmtExprVisitor,
+    PyStmtExprMutator,
+    PyStmtExprVisitorWithAnalyzer,
+    PyStmtExprMutatorWithAnalyzer,
+)
 
 # Compiler-only submodules. Skip under `TVM_USE_RUNTIME_LIB=1` since they
 # perform compiler-side FFI at module load (schema engine looks up