You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/15 14:14:09 UTC
(tvm) 04/33: [REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
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 3fdcc3f579de45d0c78caa9c873ce147861fbdbb
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:33:58 2026 +0000
[REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
---
.../contrib/test_android/test_meta_schedule.py | 2 +-
tests/python/contrib/test_example_npu.py | 2 +-
.../python/contrib/test_tir_triton_integration.py | 4 +-
tests/python/disco/test_callback.py | 4 +-
tests/python/disco/test_ccl.py | 2 +-
tests/python/disco/test_loader.py | 2 +-
tests/python/disco/test_nvshmem.py | 4 +-
tests/python/disco/test_session.py | 4 +-
tests/python/driver/test_compile.py | 4 +-
tests/python/relax/backend/adreno/mod_utils.py | 4 +-
tests/python/relax/backend/adreno/test_clml_ops.py | 2 +-
.../relax/backend/adreno/test_texture_network.py | 2 +-
.../test_transform_fold_vdevice_scope_change.py | 2 +-
.../test_distributed_transform_lower_distir.py | 2 +-
...ributed_transform_lower_global_to_local_view.py | 2 +-
...est_distributed_transform_propagate_sharding.py | 2 +-
.../relax/script/test_relax_script_basic_usage.py | 12 +--
.../relax/script/test_relax_script_distributed.py | 8 +-
.../script/test_relax_script_dynamic_shape.py | 24 +++---
.../script/test_relax_script_error_handling.py | 8 +-
.../relax/script/test_relax_script_ir_builder.py | 4 +-
.../script/test_relax_script_meta_programming.py | 8 +-
.../test_relax_script_operators_arith_cmp.py | 4 +-
.../script/test_relax_script_operators_create.py | 4 +-
.../script/test_relax_script_operators_datatype.py | 4 +-
.../script/test_relax_script_operators_grad.py | 4 +-
.../script/test_relax_script_operators_image.py | 4 +-
.../script/test_relax_script_operators_index.py | 4 +-
.../test_relax_script_operators_linear_algebra.py | 4 +-
.../test_relax_script_operators_manipulate.py | 4 +-
.../relax/script/test_relax_script_operators_nn.py | 4 +-
.../script/test_relax_script_operators_search.py | 4 +-
.../script/test_relax_script_operators_set.py | 4 +-
.../script/test_relax_script_operators_sort.py | 4 +-
.../test_relax_script_operators_statistical.py | 4 +-
.../script/test_relax_script_operators_vision.py | 4 +-
.../relax/script/test_relax_script_printer.py | 60 ++++++-------
.../relax/script/test_relax_script_pyfunc.py | 4 +-
tests/python/relax/test_analysis.py | 4 +-
.../test_analysis_computable_at_compile_time.py | 2 +-
.../relax/test_analysis_contains_impure_call.py | 2 +-
.../python/relax/test_analysis_detect_recursion.py | 4 +-
.../relax/test_analysis_estimate_memory_usage.py | 4 +-
.../test_analysis_suggest_layout_transforms.py | 2 +-
tests/python/relax/test_analysis_type_analysis.py | 2 +-
tests/python/relax/test_analysis_well_formed.py | 4 +-
tests/python/relax/test_ast_printer.py | 4 +-
.../python/relax/test_backend_dispatch_sampling.py | 4 +-
.../relax/test_backend_dispatch_sort_scan.py | 2 +-
.../relax/test_backend_transform_shape_lower.py | 4 +-
tests/python/relax/test_base_py_module_printer.py | 4 +-
.../relax/test_base_py_module_symbolic_shape.py | 4 +-
tests/python/relax/test_bind_params.py | 2 +-
tests/python/relax/test_bind_symbolic_vars.py | 2 +-
tests/python/relax/test_binding_rewrite.py | 2 +-
tests/python/relax/test_blockbuilder_core.py | 2 +-
tests/python/relax/test_blockbuilder_emit_te.py | 2 +-
tests/python/relax/test_codegen_cublas.py | 2 +-
tests/python/relax/test_codegen_cudnn.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 4 +-
tests/python/relax/test_codegen_dnnl.py | 2 +-
tests/python/relax/test_codegen_hipblas.py | 2 +-
tests/python/relax/test_codegen_tensorrt.py | 2 +-
tests/python/relax/test_contrib_vllm.py | 2 +-
tests/python/relax/test_dataflow_inplace.py | 2 +-
tests/python/relax/test_dataflow_pattern.py | 4 +-
tests/python/relax/test_dataflow_rewriter.py | 4 +-
tests/python/relax/test_dlpack_integration.py | 4 +-
tests/python/relax/test_e2e_op_dynamic.py | 2 +-
tests/python/relax/test_expr.py | 2 +-
tests/python/relax/test_expr_functor.py | 2 +-
tests/python/relax/test_fast_math_transform.py | 2 +-
tests/python/relax/test_frontend_common.py | 2 +-
tests/python/relax/test_frontend_dynamo.py | 4 +-
.../relax/test_frontend_from_exported_program.py | 2 +-
tests/python/relax/test_frontend_from_fx.py | 2 +-
tests/python/relax/test_frontend_nn_exporter.py | 2 +-
.../python/relax/test_frontend_nn_extern_module.py | 2 +-
tests/python/relax/test_frontend_nn_modules.py | 2 +-
tests/python/relax/test_frontend_nn_op.py | 4 +-
tests/python/relax/test_frontend_nn_packing.py | 2 +-
tests/python/relax/test_frontend_nn_subroutines.py | 2 +-
tests/python/relax/test_frontend_nn_tensor.py | 2 +-
tests/python/relax/test_frontend_onnx.py | 4 +-
tests/python/relax/test_frontend_stablehlo.py | 2 +-
tests/python/relax/test_inline_functions.py | 2 +-
tests/python/relax/test_kill_after_last_use.py | 2 +-
tests/python/relax/test_lower_alloc_tensor.py | 2 +-
.../relax/test_meta_schedule_relax_integration.py | 2 +-
tests/python/relax/test_op_binary.py | 2 +-
tests/python/relax/test_op_ccl.py | 2 +-
tests/python/relax/test_op_create.py | 2 +-
tests/python/relax/test_op_datatype.py | 2 +-
tests/python/relax/test_op_grad.py | 2 +-
tests/python/relax/test_op_image.py | 2 +-
tests/python/relax/test_op_index.py | 4 +-
tests/python/relax/test_op_inspect.py | 2 +-
tests/python/relax/test_op_linear_algebra.py | 2 +-
tests/python/relax/test_op_manipulate.py | 2 +-
tests/python/relax/test_op_misc.py | 4 +-
tests/python/relax/test_op_nn.py | 2 +-
tests/python/relax/test_op_nn_convolution.py | 2 +-
tests/python/relax/test_op_nn_pooling.py | 2 +-
tests/python/relax/test_op_qdq.py | 2 +-
tests/python/relax/test_op_sampling.py | 2 +-
tests/python/relax/test_op_search.py | 2 +-
tests/python/relax/test_op_set.py | 2 +-
tests/python/relax/test_op_size.py | 2 +-
tests/python/relax/test_op_sort.py | 2 +-
tests/python/relax/test_op_statistical.py | 2 +-
tests/python/relax/test_op_take.py | 2 +-
tests/python/relax/test_op_ternary.py | 2 +-
tests/python/relax/test_op_unary.py | 2 +-
tests/python/relax/test_op_view.py | 2 +-
tests/python/relax/test_op_vision.py | 2 +-
.../python/relax/test_optimize_layout_transform.py | 4 +-
tests/python/relax/test_pipeline.py | 2 +-
tests/python/relax/test_pytorch_integration.py | 4 +-
tests/python/relax/test_relax_operators.py | 2 +-
.../python/relax/test_relax_to_pyfunc_converter.py | 4 +-
.../python/relax/test_remove_redundant_reshape.py | 2 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 2 +-
tests/python/relax/test_testing_nn.py | 2 +-
tests/python/relax/test_tir_call_source_kernel.py | 4 +-
tests/python/relax/test_training_append_loss.py | 2 +-
tests/python/relax/test_training_loss.py | 2 +-
tests/python/relax/test_training_setup_trainer.py | 2 +-
.../python/relax/test_training_trainer_numeric.py | 2 +-
tests/python/relax/test_transform.py | 4 +-
.../relax/test_transform_allocate_workspace.py | 2 +-
tests/python/relax/test_transform_alter_op_impl.py | 4 +-
.../test_transform_annotate_tir_op_pattern.py | 2 +-
...st_transform_attach_attr_layout_free_tensors.py | 4 +-
.../relax/test_transform_attach_global_symbol.py | 4 +-
tests/python/relax/test_transform_bind_params.py | 4 +-
.../relax/test_transform_bind_symbolic_vars.py | 2 +-
.../relax/test_transform_bundle_model_params.py | 2 +-
.../relax/test_transform_canonicalize_bindings.py | 4 +-
tests/python/relax/test_transform_codegen_pass.py | 4 +-
.../test_transform_combine_parallel_matmul.py | 2 +-
.../relax/test_transform_compute_prim_value.py | 4 +-
.../relax/test_transform_convert_dataflow.py | 2 +-
tests/python/relax/test_transform_cse.py | 2 +-
.../relax/test_transform_dead_code_elimination.py | 2 +-
tests/python/relax/test_transform_decompose_ops.py | 2 +-
.../relax/test_transform_error_enrichment.py | 2 +-
.../relax/test_transform_expand_matmul_of_sum.py | 2 +-
.../relax/test_transform_expand_tuple_args.py | 2 +-
.../test_transform_fold_batch_norm_to_conv2d.py | 2 +-
tests/python/relax/test_transform_fold_constant.py | 4 +-
tests/python/relax/test_transform_fuse_ops.py | 4 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 4 +-
tests/python/relax/test_transform_fuse_tir.py | 4 +-
.../relax/test_transform_fuse_transpose_matmul.py | 4 +-
tests/python/relax/test_transform_gradient.py | 2 +-
.../relax/test_transform_gradient_te_register.py | 2 +-
.../test_transform_inline_private_functions.py | 2 +-
.../relax/test_transform_ipc_allreduce_rewrite.py | 2 +-
tests/python/relax/test_transform_lambda_lift.py | 4 +-
.../relax/test_transform_lazy_transform_params.py | 4 +-
tests/python/relax/test_transform_legalize_ops.py | 4 +-
.../relax/test_transform_legalize_ops_binary.py | 4 +-
.../relax/test_transform_legalize_ops_ccl.py | 4 +-
.../test_transform_legalize_ops_create_datatype.py | 4 +-
.../test_transform_legalize_ops_distributed.py | 4 +-
.../relax/test_transform_legalize_ops_grad.py | 4 +-
.../relax/test_transform_legalize_ops_image.py | 4 +-
..._transform_legalize_ops_index_linear_algebra.py | 4 +-
.../test_transform_legalize_ops_manipulate.py | 4 +-
.../python/relax/test_transform_legalize_ops_nn.py | 4 +-
.../relax/test_transform_legalize_ops_qdq.py | 4 +-
...st_transform_legalize_ops_search_statistical.py | 4 +-
.../relax/test_transform_legalize_ops_unary.py | 2 +-
.../relax/test_transform_lift_transform_params.py | 4 +-
.../test_transform_lower_gpu_ipc_alloc_storage.py | 2 +-
.../test_transform_merge_composite_functions.py | 4 +-
.../test_transform_meta_schedule_apply_database.py | 2 +-
.../relax/test_transform_meta_schedule_tuning.py | 4 +-
tests/python/relax/test_transform_normalize.py | 2 +-
.../relax/test_transform_normalize_global_var.py | 4 +-
...st_transform_operator_specific_normalization.py | 2 +-
.../relax/test_transform_remove_unused_outputs.py | 2 +-
.../test_transform_remove_unused_parameters.py | 2 +-
..._transform_reorder_permute_dims_after_concat.py | 2 +-
.../test_transform_reorder_take_after_matmul.py | 2 +-
.../relax/test_transform_rewrite_cuda_graph.py | 4 +-
.../test_transform_rewrite_dataflow_reshape.py | 4 +-
...nsform_specialize_function_based_on_callsite.py | 2 +-
.../test_transform_split_layout_rewrite_preproc.py | 4 +-
.../test_transform_static_plan_block_memory.py | 4 +-
.../relax/test_transform_to_mixed_precision.py | 2 +-
.../relax/test_transform_topological_sort.py | 2 +-
.../relax/test_transform_update_param_type.py | 2 +-
.../relax/test_vm_alloc_storage_with_scope.py | 4 +-
tests/python/relax/test_vm_build.py | 4 +-
tests/python/relax/test_vm_builtin.py | 2 +-
tests/python/relax/test_vm_builtin_lower.py | 2 +-
tests/python/relax/test_vm_callback_function.py | 2 +-
tests/python/relax/test_vm_codegen_only.py | 4 +-
tests/python/relax/test_vm_cuda_graph.py | 4 +-
tests/python/relax/test_vm_execbuilder.py | 2 +-
tests/python/relax/texture/test_texture_nd.py | 2 +-
...st_s_tir_analysis_calculate_allocated_memory.py | 2 +-
.../test_s_tir_analysis_estimate_tir_flops.py | 2 +-
.../test_s_tir_analysis_identify_memcpy.py | 2 +-
.../test_s_tir_analysis_is_pure_function.py | 2 +-
.../s_tir/analysis/test_s_tir_analysis_oob.py | 2 +-
.../test_s_tir_analysis_verify_block_scope.py | 2 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../analysis/test_sblock_tensor_access_lca.py | 2 +-
.../python/s_tir/base/test_compilation_pipeline.py | 2 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 2 +-
.../s_tir/base/test_sblock_dependence_info.py | 2 +-
.../s_tir/base/test_tir_te_extern_function.py | 2 +-
tests/python/s_tir/dlight/test_benchmark.py | 4 +-
tests/python/s_tir/dlight/test_cpu_gemv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_conv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_fallback.py | 2 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 2 +-
.../s_tir/dlight/test_gpu_general_reduction.py | 2 +-
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_matmul.py | 2 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 2 +-
tests/python/s_tir/dlight/test_gpu_reduction.py | 2 +-
tests/python/s_tir/dlight/test_gpu_rmsnorm.py | 2 +-
tests/python/s_tir/dlight/test_gpu_transpose.py | 2 +-
tests/python/s_tir/dlight/test_primitives.py | 2 +-
.../meta_schedule/test_meta_schedule_arg_info.py | 2 +-
.../meta_schedule/test_meta_schedule_builder.py | 2 +-
.../meta_schedule/test_meta_schedule_cost_model.py | 2 +-
.../meta_schedule/test_meta_schedule_database.py | 2 +-
...schedule_feature_extractor_per_store_feature.py | 2 +-
.../test_meta_schedule_measure_callback.py | 2 +-
...eta_schedule_mutator_mutate_compute_location.py | 2 +-
.../test_meta_schedule_mutator_mutate_parallel.py | 2 +-
..._meta_schedule_mutator_mutate_thread_binding.py | 2 +-
.../test_meta_schedule_mutator_mutate_tile_size.py | 2 +-
.../test_meta_schedule_mutator_mutate_unroll.py | 2 +-
.../test_meta_schedule_post_order_apply.py | 2 +-
...ule_postproc_disallow_async_strided_mem_copy.py | 2 +-
...meta_schedule_postproc_disallow_dynamic_loop.py | 2 +-
..._schedule_postproc_rewrite_cooperative_fetch.py | 2 +-
.../test_meta_schedule_postproc_rewrite_layout.py | 2 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 2 +-
...ta_schedule_postproc_rewrite_reduction_block.py | 2 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 2 +-
...meta_schedule_postproc_rewrite_unbound_block.py | 2 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 2 +-
...est_meta_schedule_postproc_verify_vtcm_limit.py | 2 +-
.../meta_schedule/test_meta_schedule_runner.py | 2 +-
...test_meta_schedule_schedule_rule_add_rfactor.py | 2 +-
...eta_schedule_schedule_rule_apply_custom_rule.py | 2 +-
.../test_meta_schedule_schedule_rule_auto_bind.py | 2 +-
...test_meta_schedule_schedule_rule_auto_inline.py | 2 +-
...chedule_schedule_rule_cross_thread_reduction.py | 2 +-
.../test_meta_schedule_schedule_rule_mlt.py | 2 +-
.../test_meta_schedule_schedule_rule_mlt_intrin.py | 2 +-
...dule_schedule_rule_parallel_vectorize_unroll.py | 2 +-
...hedule_schedule_rule_random_compute_location.py | 2 +-
.../test_meta_schedule_search_strategy.py | 2 +-
.../meta_schedule/test_meta_schedule_space_cpu.py | 2 +-
.../meta_schedule/test_meta_schedule_space_cuda.py | 2 +-
.../test_meta_schedule_space_cuda_async.py | 2 +-
.../test_meta_schedule_space_generator.py | 2 +-
.../test_meta_schedule_space_post_opt.py | 2 +-
.../test_meta_schedule_task_scheduler.py | 2 +-
.../test_meta_schedule_trace_apply.py | 2 +-
.../test_meta_schedule_tune_context.py | 2 +-
.../meta_schedule/test_meta_schedule_tune_tir.py | 2 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 2 +-
.../test_tir_schedule_annotate_tensor_access.py | 2 +-
.../schedule/test_tir_schedule_block_scope.py | 2 +-
.../s_tir/schedule/test_tir_schedule_blockize.py | 2 +-
.../schedule/test_tir_schedule_cache_index.py | 2 +-
.../schedule/test_tir_schedule_cache_read_write.py | 2 +-
.../s_tir/schedule/test_tir_schedule_compute_at.py | 2 +-
.../schedule/test_tir_schedule_compute_inline.py | 2 +-
.../test_tir_schedule_decompose_padding.py | 2 +-
.../s_tir/schedule/test_tir_schedule_error.py | 2 +-
.../s_tir/schedule/test_tir_schedule_for_kind.py | 2 +-
.../test_tir_schedule_fuse_reduction_epilogue.py | 2 +-
...ir_schedule_fuse_reduction_epilogue_clipping.py | 2 +-
...st_tir_schedule_fuse_reduction_epilogue_relu.py | 2 +-
.../s_tir/schedule/test_tir_schedule_merge.py | 2 +-
.../s_tir/schedule/test_tir_schedule_pad_einsum.py | 2 +-
.../s_tir/schedule/test_tir_schedule_partition.py | 2 +-
.../schedule/test_tir_schedule_read_write_at.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reduction.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reindex.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reorder.py | 2 +-
.../test_tir_schedule_reorder_block_iter_var.py | 2 +-
.../s_tir/schedule/test_tir_schedule_rfactor.py | 2 +-
.../schedule/test_tir_schedule_rolling_buffer.py | 2 +-
.../s_tir/schedule/test_tir_schedule_sampling.py | 2 +-
.../s_tir/schedule/test_tir_schedule_set_dtype.py | 2 +-
.../s_tir/schedule/test_tir_schedule_set_scope.py | 2 +-
.../s_tir/schedule/test_tir_schedule_split_fuse.py | 2 +-
.../s_tir/schedule/test_tir_schedule_state.py | 10 +--
.../test_tir_schedule_state_cached_flags.py | 2 +-
.../schedule/test_tir_schedule_storage_align.py | 2 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 2 +-
.../s_tir/schedule/test_tir_schedule_trace.py | 2 +-
.../s_tir/schedule/test_tir_schedule_transform.py | 2 +-
.../schedule/test_tir_schedule_transform_layout.py | 2 +-
.../s_tir/schedule/test_tir_schedule_utilities.py | 2 +-
.../schedule/test_tir_unsafe_hide_tensor_access.py | 2 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 16 ++--
.../s_tir/script/test_s_tir_script_blocks.py | 40 ++++-----
.../script/test_s_tir_script_dynamic_shape.py | 8 +-
.../script/test_s_tir_script_error_handling.py | 2 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 2 +-
.../script/test_s_tir_script_meta_programming.py | 4 +-
.../s_tir/script/test_s_tir_script_namespace.py | 22 ++---
.../s_tir/script/test_s_tir_script_printer.py | 82 +++++++++---------
.../script/test_s_tir_script_source_locations.py | 2 +-
tests/python/s_tir/test_arith_domain_touched.py | 2 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 2 +-
...test_s_tir_transform_annotate_irregular_loop.py | 2 +-
.../test_s_tir_transform_canonicalize_loop.py | 2 +-
.../test_s_tir_transform_compact_tensor_region.py | 2 +-
...est_s_tir_transform_convert_blocks_to_opaque.py | 2 +-
.../test_s_tir_transform_default_gpu_schedule.py | 2 +-
.../test_s_tir_transform_inject_permuted_layout.py | 2 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 4 +-
...est_s_tir_transform_inject_software_pipeline.py | 2 +-
.../test_s_tir_transform_inject_virtual_thread.py | 2 +-
.../test_s_tir_transform_lift_thread_binding.py | 2 +-
.../test_s_tir_transform_loop_partition.py | 2 +-
...s_tir_transform_lower_cross_thread_reduction.py | 2 +-
.../test_s_tir_transform_lower_init_block.py | 2 +-
.../test_s_tir_transform_lower_match_tensor.py | 2 +-
.../test_s_tir_transform_lower_opaque_block.py | 2 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 2 +-
.../test_s_tir_transform_lower_thread_binding.py | 2 +-
...transform_manifest_shared_memory_local_stage.py | 2 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 2 +-
...form_merge_dynamic_shared_memory_allocations.py | 2 +-
.../test_s_tir_transform_mma_tensor_layout.py | 2 +-
...sform_plan_update_tensor_allocation_location.py | 2 +-
.../transform/test_s_tir_transform_remove_undef.py | 2 +-
...transform_remove_weight_layout_rewrite_block.py | 2 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 2 +-
.../transform/test_s_tir_transform_thread_sync.py | 2 +-
tests/python/script/test_constructor_contracts.py | 2 +-
tests/python/target/test_arm_target.py | 2 +-
tests/python/target/test_target_target.py | 2 +-
tests/python/te/test_te_create_function.py | 4 +-
.../codegen/conftest.py} | 34 +++-----
tests/python/tirx/codegen/test_tirx_codegen_fp8.py | 2 +-
.../conftest.py} | 28 +------
.../tvm_tirx/testing/_pytest.py} | 34 ++++++++
.../python/runtime/test_runtime_memory_scope.py | 2 +-
tirx/tests/python/tirx/codegen/conftest.py | 43 +++-------
tirx/tests/python/tirx/conftest.py | 97 +---------------------
.../python/tirx/script/test_tirx_script_tracing.py | 2 +-
tirx/tests/python/tirx/test_cuda_arch_marker.py | 2 +-
tirx/tests/python/tirx/test_cuda_backend_config.py | 2 +-
.../transform/test_tirx_transform_lower_intrin.py | 4 +-
358 files changed, 642 insertions(+), 750 deletions(-)
diff --git a/tests/python/contrib/test_android/test_meta_schedule.py b/tests/python/contrib/test_android/test_meta_schedule.py
index 4d28c6d786..1eab8b2b13 100644
--- a/tests/python/contrib/test_android/test_meta_schedule.py
+++ b/tests/python/contrib/test_android/test_meta_schedule.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
import tvm.topi.testing
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.builder import LocalBuilder
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from .infrastructure import get_android_gpu_target, get_rpc_runner
diff --git a/tests/python/contrib/test_example_npu.py b/tests/python/contrib/test_example_npu.py
index b0cee54f5a..9617e7a745 100644
--- a/tests/python/contrib/test_example_npu.py
+++ b/tests/python/contrib/test_example_npu.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/contrib/test_tir_triton_integration.py b/tests/python/contrib/test_tir_triton_integration.py
index d4532f570a..ce5fb0122e 100644
--- a/tests/python/contrib/test_tir_triton_integration.py
+++ b/tests/python/contrib/test_tir_triton_integration.py
@@ -25,8 +25,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_callback.py b/tests/python/disco/test_callback.py
index 55d362017e..403cdc9c89 100644
--- a/tests/python/disco/test_callback.py
+++ b/tests/python/disco/test_callback.py
@@ -26,8 +26,8 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_ccl.py b/tests/python/disco/test_ccl.py
index 56ef4b2791..64d091de3d 100644
--- a/tests/python/disco/test_ccl.py
+++ b/tests/python/disco/test_ccl.py
@@ -29,7 +29,7 @@ from tvm import relax as rx
from tvm.runtime import disco as di
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
if di is None:
pytest.skip("disco runtime is not available", allow_module_level=True)
diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py
index f4299df4b1..017239a1e4 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -32,7 +32,7 @@ from tvm.contrib import tvmjs
from tvm.runtime import disco as di
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_nvshmem.py b/tests/python/disco/test_nvshmem.py
index d0b1ca58c3..2b8b4a8909 100644
--- a/tests/python/disco/test_nvshmem.py
+++ b/tests/python/disco/test_nvshmem.py
@@ -34,8 +34,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.runtime import disco as di
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py
index 61c672804c..375c2d3d6b 100644
--- a/tests/python/disco/test_session.py
+++ b/tests/python/disco/test_session.py
@@ -35,8 +35,8 @@ import tvm_tirx.testing, tvm
from tvm.exec import disco_worker as _ # noqa: F401 # pylint: disable=unused-import
from tvm.runtime import disco as di
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
if di is None:
diff --git a/tests/python/driver/test_compile.py b/tests/python/driver/test_compile.py
index cba770948a..b71ba67b0a 100644
--- a/tests/python/driver/test_compile.py
+++ b/tests/python/driver/test_compile.py
@@ -25,8 +25,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, te
from tvm_tirx.runtime import Executable
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/backend/adreno/mod_utils.py b/tests/python/relax/backend/adreno/mod_utils.py
index f239adf24e..fd3e2fe4a8 100644
--- a/tests/python/relax/backend/adreno/mod_utils.py
+++ b/tests/python/relax/backend/adreno/mod_utils.py
@@ -28,8 +28,8 @@ from tvm import relax, rpc
from tvm.relax.backend.adreno import clml
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_clml_ops.py b/tests/python/relax/backend/adreno/test_clml_ops.py
index 45cb5e775f..08ade82ee7 100644
--- a/tests/python/relax/backend/adreno/test_clml_ops.py
+++ b/tests/python/relax/backend/adreno/test_clml_ops.py
@@ -53,7 +53,7 @@ from tvm.relax.backend.adreno import clml
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad, OpenCLMLOffLoadForLLM
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_texture_network.py b/tests/python/relax/backend/adreno/test_texture_network.py
index aeeb39ac76..8039372deb 100644
--- a/tests/python/relax/backend/adreno/test_texture_network.py
+++ b/tests/python/relax/backend/adreno/test_texture_network.py
@@ -34,7 +34,7 @@ from tvm.relax.frontend.onnx import from_onnx
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
index bec2900153..1fa999d25d 100644
--- a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
+++ b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
index fadbd12c3a..f5fe112c3f 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
@@ -24,7 +24,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
index a121224e07..891e201d7f 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
index 4963526601..3ec43750ed 100644
--- a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
+++ b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
@@ -25,7 +25,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_basic_usage.py b/tests/python/relax/script/test_relax_script_basic_usage.py
index b28d9f0240..8f631bd0ee 100644
--- a/tests/python/relax/script/test_relax_script_basic_usage.py
+++ b/tests/python/relax/script/test_relax_script_basic_usage.py
@@ -32,8 +32,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import DummyGlobalInfo, VDevice
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -46,9 +46,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
@@ -1550,8 +1550,8 @@ def test_roundtrip_basic_usage(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_distributed.py b/tests/python/relax/script/test_relax_script_distributed.py
index 85177dc543..009ebe93d9 100644
--- a/tests/python/relax/script/test_relax_script_distributed.py
+++ b/tests/python/relax/script/test_relax_script_distributed.py
@@ -27,8 +27,8 @@ from tvm_tirx.ir import Range
from tvm.relax import Call, SeqExpr, VarBinding
from tvm.relax.distributed import DeviceMesh
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -41,9 +41,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py
index 72abb44241..4d1760f972 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -30,8 +30,8 @@ from tvm_tirx import IRModule, tirx
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
M = T.dynamic("M", "int64")
@@ -59,7 +59,7 @@ def test_type_vars_roundtrip():
def func[M: T.int64](x: R.Tensor((M, M * 2), "float32")):
return x
""",
- extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx},
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
)
tvm_tirx.ir.assert_structural_equal(func, typed)
assert 'M = T.dynamic("M", dtype="int64")' not in script
@@ -76,13 +76,13 @@ def func[M: T.int64](x: R.Tensor((M, M * 2), "float32")):
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ script, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
),
)
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- repeated, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ repeated, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
),
)
@@ -105,7 +105,7 @@ def test_dynamic_module_symbol_identity():
restored = tvm_tirx.script.from_source(
source,
check_well_formed=False,
- extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx},
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
)
first_n = restored["first"].params[0].ty.shape.values[0]
second_shape = restored["second"].params[0].ty.shape.values
@@ -124,9 +124,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
@@ -250,9 +250,9 @@ def test_call_tir_with_tir_var():
repeated,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
)
@@ -651,8 +651,8 @@ def test_roundtrip_dynamic_shape(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_error_handling.py b/tests/python/relax/script/test_relax_script_error_handling.py
index da44aa11c3..e32592fe51 100644
--- a/tests/python/relax/script/test_relax_script_error_handling.py
+++ b/tests/python/relax/script/test_relax_script_error_handling.py
@@ -31,8 +31,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -45,9 +45,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_ir_builder.py b/tests/python/relax/script/test_relax_script_ir_builder.py
index 9121507813..4ec4557713 100644
--- a/tests/python/relax/script/test_relax_script_ir_builder.py
+++ b/tests/python/relax/script/test_relax_script_ir_builder.py
@@ -193,7 +193,7 @@ def test_function_subroutine_before_main():
"""The block builder can generate subroutines, and calls into subroutines"""
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
# create with TVMScript
@I.ir_module
@@ -235,7 +235,7 @@ def test_function_subroutine_during_main():
"""Subroutines may be generated as needed, pausing the main function collection"""
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
# create with TVMScript
@I.ir_module
diff --git a/tests/python/relax/script/test_relax_script_meta_programming.py b/tests/python/relax/script/test_relax_script_meta_programming.py
index 006d4435d9..557514d830 100644
--- a/tests/python/relax/script/test_relax_script_meta_programming.py
+++ b/tests/python/relax/script/test_relax_script_meta_programming.py
@@ -30,8 +30,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx import IRModule, ir
from tvm import relax, topi
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -44,9 +44,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
index b0228601e4..3a14bbb24e 100644
--- a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
+++ b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
@@ -25,7 +25,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -34,7 +34,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_create.py b/tests/python/relax/script/test_relax_script_operators_create.py
index 89adacd68a..043ecf19e1 100644
--- a/tests/python/relax/script/test_relax_script_operators_create.py
+++ b/tests/python/relax/script/test_relax_script_operators_create.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_datatype.py b/tests/python/relax/script/test_relax_script_operators_datatype.py
index 2026ff5351..f394203bf4 100644
--- a/tests/python/relax/script/test_relax_script_operators_datatype.py
+++ b/tests/python/relax/script/test_relax_script_operators_datatype.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_grad.py b/tests/python/relax/script/test_relax_script_operators_grad.py
index 62fd09d7a4..bd43f81248 100644
--- a/tests/python/relax/script/test_relax_script_operators_grad.py
+++ b/tests/python/relax/script/test_relax_script_operators_grad.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -29,7 +29,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_image.py b/tests/python/relax/script/test_relax_script_operators_image.py
index d59618d804..53cd473565 100644
--- a/tests/python/relax/script/test_relax_script_operators_image.py
+++ b/tests/python/relax/script/test_relax_script_operators_image.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_index.py b/tests/python/relax/script/test_relax_script_operators_index.py
index ea96339137..63bb8aa7c8 100644
--- a/tests/python/relax/script/test_relax_script_operators_index.py
+++ b/tests/python/relax/script/test_relax_script_operators_index.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
index f5d33ebdff..0d7f37000d 100644
--- a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
+++ b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_manipulate.py b/tests/python/relax/script/test_relax_script_operators_manipulate.py
index 8fa3d5b380..5119f8f6cc 100644
--- a/tests/python/relax/script/test_relax_script_operators_manipulate.py
+++ b/tests/python/relax/script/test_relax_script_operators_manipulate.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -31,7 +31,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_nn.py b/tests/python/relax/script/test_relax_script_operators_nn.py
index a81e09ae46..4268d3b7f3 100644
--- a/tests/python/relax/script/test_relax_script_operators_nn.py
+++ b/tests/python/relax/script/test_relax_script_operators_nn.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_search.py b/tests/python/relax/script/test_relax_script_operators_search.py
index ba99c0494a..9da26aee1d 100644
--- a/tests/python/relax/script/test_relax_script_operators_search.py
+++ b/tests/python/relax/script/test_relax_script_operators_search.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_set.py b/tests/python/relax/script/test_relax_script_operators_set.py
index 44f509873d..7781de2e9c 100644
--- a/tests/python/relax/script/test_relax_script_operators_set.py
+++ b/tests/python/relax/script/test_relax_script_operators_set.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_sort.py b/tests/python/relax/script/test_relax_script_operators_sort.py
index bd263acd70..15cf8f6f48 100644
--- a/tests/python/relax/script/test_relax_script_operators_sort.py
+++ b/tests/python/relax/script/test_relax_script_operators_sort.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_statistical.py b/tests/python/relax/script/test_relax_script_operators_statistical.py
index 12bdd5dcfd..569c230b06 100644
--- a/tests/python/relax/script/test_relax_script_operators_statistical.py
+++ b/tests/python/relax/script/test_relax_script_operators_statistical.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_vision.py b/tests/python/relax/script/test_relax_script_operators_vision.py
index e9ae12fa83..666ee0a0de 100644
--- a/tests/python/relax/script/test_relax_script_operators_vision.py
+++ b/tests/python/relax/script/test_relax_script_operators_vision.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index e690435747..d6abc532ed 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -32,8 +32,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import TensorType
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.printer.scriptable import PrinterConfig, _script
@@ -166,7 +166,7 @@ def test_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@R.function
@@ -183,8 +183,8 @@ def test_module():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -234,7 +234,7 @@ def test_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
@@ -315,7 +315,7 @@ def test_lone_private_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function(private=True)
def main(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
@@ -340,7 +340,7 @@ def test_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -366,7 +366,7 @@ def test_extern_func_with_ty():
_assert_print_lines(
obj,
"""
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -389,9 +389,9 @@ def test_extern_func_with_ty_roundtrip():
mod.script(verbose_expr=True),
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(mod, roundtrip)
@@ -420,7 +420,7 @@ def test_nested_function(relax_prefix, comment_imports):
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -581,7 +581,7 @@ def test_prim_value():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@R.function
@@ -601,9 +601,9 @@ def func() -> T.int64:
float_script,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
float_func,
@@ -627,9 +627,9 @@ def test_primitive_bindings_roundtrip_without_prim_value_marker():
source,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
func,
@@ -959,8 +959,8 @@ def test_module_cross_func_call():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -983,8 +983,8 @@ class Module:
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1014,7 +1014,7 @@ def test_assert_op():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1040,7 +1040,7 @@ def test_print():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1066,7 +1066,7 @@ def test_private_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1121,7 +1121,7 @@ def test_directly_construct_private_funcs():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1168,7 +1168,7 @@ def test_reused_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
@@ -1198,7 +1198,7 @@ def test_inline_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
@@ -1248,7 +1248,7 @@ def test_hide_inferable_ty():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(A: R.Tensor((10, 20), dtype="float32"), B: R.Tensor(dtype="float32", ndim=2)) -> R.Tensor((10, 20), dtype="float32"):
@@ -1286,8 +1286,8 @@ def test_extern_func_roundtrip(show_all_ty):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_pyfunc.py b/tests/python/relax/script/test_relax_script_pyfunc.py
index 8d298f30c0..38cfd0ed8f 100644
--- a/tests/python/relax/script/test_relax_script_pyfunc.py
+++ b/tests/python/relax/script/test_relax_script_pyfunc.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py
index ac5fe5aece..b6f4edf585 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -34,8 +34,8 @@ from tvm.relax.analysis import (
used_vars,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_computable_at_compile_time.py b/tests/python/relax/test_analysis_computable_at_compile_time.py
index 5a9f496aca..f44ac42d0a 100644
--- a/tests/python/relax/test_analysis_computable_at_compile_time.py
+++ b/tests/python/relax/test_analysis_computable_at_compile_time.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_contains_impure_call.py b/tests/python/relax/test_analysis_contains_impure_call.py
index 1f3d2158cd..a721f4f595 100644
--- a/tests/python/relax/test_analysis_contains_impure_call.py
+++ b/tests/python/relax/test_analysis_contains_impure_call.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
from tvm.relax.analysis import contains_impure_call
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple_pure_case():
diff --git a/tests/python/relax/test_analysis_detect_recursion.py b/tests/python/relax/test_analysis_detect_recursion.py
index f90f0bfe1c..b2ca5889c6 100644
--- a/tests/python/relax/test_analysis_detect_recursion.py
+++ b/tests/python/relax/test_analysis_detect_recursion.py
@@ -19,8 +19,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
from tvm.relax.analysis import detect_recursion
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_estimate_memory_usage.py b/tests/python/relax/test_analysis_estimate_memory_usage.py
index 7396cc3bc1..615053ac15 100644
--- a/tests/python/relax/test_analysis_estimate_memory_usage.py
+++ b/tests/python/relax/test_analysis_estimate_memory_usage.py
@@ -19,8 +19,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.analysis import estimate_memory_usage
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py
index cb873f83fc..3ead1c5b91 100644
--- a/tests/python/relax/test_analysis_suggest_layout_transforms.py
+++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py
@@ -21,7 +21,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax, s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_type_analysis.py b/tests/python/relax/test_analysis_type_analysis.py
index 2dcf60f998..e1a3b9ea9e 100644
--- a/tests/python/relax/test_analysis_type_analysis.py
+++ b/tests/python/relax/test_analysis_type_analysis.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import ir, tirx
from tvm import relax as rx
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 4ca5ae5ff8..1f7ed61172 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -27,8 +27,8 @@ from tvm_tirx import tirx
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
diff --git a/tests/python/relax/test_ast_printer.py b/tests/python/relax/test_ast_printer.py
index cb3be30e60..e784134a05 100644
--- a/tests/python/relax/test_ast_printer.py
+++ b/tests/python/relax/test_ast_printer.py
@@ -27,8 +27,8 @@ from tvm_tirx import tirx
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.testing import dump_ast
from tvm.relax.testing.ast_printer import ASTPrinter
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# Overload dump_ast to test both type and type annotations
diff --git a/tests/python/relax/test_backend_dispatch_sampling.py b/tests/python/relax/test_backend_dispatch_sampling.py
index df5424f331..53cce1dea3 100644
--- a/tests/python/relax/test_backend_dispatch_sampling.py
+++ b/tests/python/relax/test_backend_dispatch_sampling.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.backend import DispatchSampling
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py
index c8446535a3..bf8709e142 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -31,7 +31,7 @@ from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.backend import DispatchSortScan
from tvm.s_tir import dlight
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py b/tests/python/relax/test_backend_transform_shape_lower.py
index 6739031b33..eaa9bb6a4f 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -22,8 +22,8 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# note: we expected RemovePurityChecking to be run first, so we force purity in most test cases
diff --git a/tests/python/relax/test_base_py_module_printer.py b/tests/python/relax/test_base_py_module_printer.py
index fa6aef20d7..71211c32ee 100644
--- a/tests/python/relax/test_base_py_module_printer.py
+++ b/tests/python/relax/test_base_py_module_printer.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
from tvm_tirx.ir import StringImm
from tvm.relax.base_py_module import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_base_py_module_symbolic_shape.py b/tests/python/relax/test_base_py_module_symbolic_shape.py
index 9301fe53d7..a5a5398cb7 100644
--- a/tests/python/relax/test_base_py_module_symbolic_shape.py
+++ b/tests/python/relax/test_base_py_module_symbolic_shape.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import IRModule
from tvm.relax.base_py_module import BasePyModule
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_bind_params.py b/tests/python/relax/test_bind_params.py
index 4d467a7060..9d6a205e35 100644
--- a/tests/python/relax/test_bind_params.py
+++ b/tests/python/relax/test_bind_params.py
@@ -24,7 +24,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
param_specification = tvm_tirx.testing.parameter("by_string", "by_var")
param_shape = tvm_tirx.testing.parameter("static_shape", "dynamic_shape", "ndim", "arbitrary")
diff --git a/tests/python/relax/test_bind_symbolic_vars.py b/tests/python/relax/test_bind_symbolic_vars.py
index e5640636db..00d2790fa1 100644
--- a/tests/python/relax/test_bind_symbolic_vars.py
+++ b/tests/python/relax/test_bind_symbolic_vars.py
@@ -19,7 +19,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
replace_by_tir_var = tvm_tirx.testing.parameter(
diff --git a/tests/python/relax/test_binding_rewrite.py b/tests/python/relax/test_binding_rewrite.py
index b8771fff53..7b37c6b908 100644
--- a/tests/python/relax/test_binding_rewrite.py
+++ b/tests/python/relax/test_binding_rewrite.py
@@ -24,7 +24,7 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.analysis import name_to_binding
from tvm.relax.binding_rewrite import DataflowBlockRewrite
from tvm.relax.expr import DataflowVar, Var
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py
index 48a11a655a..6dee48dd1b 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -29,7 +29,7 @@ from tvm import te, topi
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/relax/test_blockbuilder_emit_te.py b/tests/python/relax/test_blockbuilder_emit_te.py
index b94671b445..08c8c7bf4c 100644
--- a/tests/python/relax/test_blockbuilder_emit_te.py
+++ b/tests/python/relax/test_blockbuilder_emit_te.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
from tvm import relax as rx
from tvm import te
from tvm_tirx.ir.base import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_codegen_cublas.py b/tests/python/relax/test_codegen_cublas.py
index 8c403286ed..f0ee137d19 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.backend.cuda.cublas import partition_for_cublas
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_matmul_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
try:
diff --git a/tests/python/relax/test_codegen_cudnn.py b/tests/python/relax/test_codegen_cudnn.py
index 407a31086a..2bf2d540ab 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -30,7 +30,7 @@ from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.cuda.cudnn import partition_for_cudnn
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_stacked_attention_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py
index a87c6bab2e..f1f50ec6e7 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -38,8 +38,8 @@ from tvm.relax.testing import (
get_relax_stacked_attention_module,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py
index 2d03350ee9..e3c7a38779 100644
--- a/tests/python/relax/test_codegen_dnnl.py
+++ b/tests/python/relax/test_codegen_dnnl.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.dpl import make_fused_bias_activation_pattern
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_codegen_hipblas.py b/tests/python/relax/test_codegen_hipblas.py
index cba5382b32..26c47fa4a1 100644
--- a/tests/python/relax/test_codegen_hipblas.py
+++ b/tests/python/relax/test_codegen_hipblas.py
@@ -27,7 +27,7 @@ import tvm.topi.testing
from tvm import relax
from tvm.relax.backend.rocm.hipblas import partition_for_hipblas
from tvm.relax.testing import get_relax_matmul_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
try:
import ml_dtypes
diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py
index 3f89e5437f..43a7d3b643 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -23,7 +23,7 @@ from tvm import relax
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.dpl import is_op, make_fused_bias_activation_pattern, wildcard
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_contrib_vllm.py b/tests/python/relax/test_contrib_vllm.py
index b1891a3064..53ee0475c0 100644
--- a/tests/python/relax/test_contrib_vllm.py
+++ b/tests/python/relax/test_contrib_vllm.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_dataflow_inplace.py b/tests/python/relax/test_dataflow_inplace.py
index e06f1680a7..0e71f2bb4a 100644
--- a/tests/python/relax/test_dataflow_inplace.py
+++ b/tests/python/relax/test_dataflow_inplace.py
@@ -31,7 +31,7 @@ from tvm.relax.testing.transform import (
dataflow_single_inplace_call,
)
from tvm.relax.transform import DataflowUseInplaceCalls
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_dataflow_pattern.py b/tests/python/relax/test_dataflow_pattern.py
index 629cb2bf59..d640bcbf65 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -31,8 +31,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.analysis import get_var2val
from tvm.relax.dpl import *
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dataflow_rewriter.py b/tests/python/relax/test_dataflow_rewriter.py
index bce28e456c..5c60778706 100644
--- a/tests/python/relax/test_dataflow_rewriter.py
+++ b/tests/python/relax/test_dataflow_rewriter.py
@@ -20,8 +20,8 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dlpack_integration.py b/tests/python/relax/test_dlpack_integration.py
index 8f4b7786df..45427c6aaa 100644
--- a/tests/python/relax/test_dlpack_integration.py
+++ b/tests/python/relax/test_dlpack_integration.py
@@ -30,8 +30,8 @@ import torch
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import BasePyModule
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_e2e_op_dynamic.py b/tests/python/relax/test_e2e_op_dynamic.py
index 3f3d1ba6e4..fb111c89a6 100644
--- a/tests/python/relax/test_e2e_op_dynamic.py
+++ b/tests/python/relax/test_e2e_op_dynamic.py
@@ -27,7 +27,7 @@ import tvm.topi.testing
from tvm_tirx import tirx
from tvm import relax
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# TODO(tvm-team): `tirx.transform.DefaultGPUSchedule` does not work.
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index e054ffab1b..964c605e17 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
from tvm import relax as rx
from tvm_tirx import tirx
from tvm.relax.expr import make_shape
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_equal(x, y, map_free_vars=False):
diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py
index cf8302db39..d8b563410a 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -41,7 +41,7 @@ from tvm.relax.expr import (
Var,
VarBinding,
)
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
m, n = tirx.Var("m", "int64"), tirx.Var("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
diff --git a/tests/python/relax/test_fast_math_transform.py b/tests/python/relax/test_fast_math_transform.py
index 8a2e6e4b88..51804983c7 100644
--- a/tests/python/relax/test_fast_math_transform.py
+++ b/tests/python/relax/test_fast_math_transform.py
@@ -24,7 +24,7 @@ from tvm import relax, topi
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform import FastMathTransform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_frontend_common.py b/tests/python/relax/test_frontend_common.py
index 8420b0efdf..d49903c5b3 100644
--- a/tests/python/relax/test_frontend_common.py
+++ b/tests/python/relax/test_frontend_common.py
@@ -20,7 +20,7 @@ from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.common import autopad
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py
index 722bd6636c..b69c5b57fc 100644
--- a/tests/python/relax/test_frontend_dynamo.py
+++ b/tests/python/relax/test_frontend_dynamo.py
@@ -31,8 +31,8 @@ from tvm import relax
from tvm.relax.frontend.torch import relax_dynamo
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py
index 87978cb483..cbe277fb11 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index 556ba02ceb..f84f7ee92c 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.torch import from_fx
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_exporter.py b/tests/python/relax/test_frontend_nn_exporter.py
index 67fc9362e6..86d34e84aa 100644
--- a/tests/python/relax/test_frontend_nn_exporter.py
+++ b/tests/python/relax/test_frontend_nn_exporter.py
@@ -24,7 +24,7 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_extern_module.py b/tests/python/relax/test_frontend_nn_extern_module.py
index a8c3c9c0f4..06cf8f9908 100644
--- a/tests/python/relax/test_frontend_nn_extern_module.py
+++ b/tests/python/relax/test_frontend_nn_extern_module.py
@@ -82,7 +82,7 @@ def _test_infer_sym(func, x, y, z): # pylint: disable=invalid-name
def _check_ir_equality(mod):
# pylint: disable=import-outside-toplevel
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# pylint: enable=import-outside-toplevel
diff --git a/tests/python/relax/test_frontend_nn_modules.py b/tests/python/relax/test_frontend_nn_modules.py
index 118914384d..ab6d1d9b52 100644
--- a/tests/python/relax/test_frontend_nn_modules.py
+++ b/tests/python/relax/test_frontend_nn_modules.py
@@ -26,7 +26,7 @@ from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import core, modules, spec
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py
index 4b52ff7f41..676692730c 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -27,8 +27,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.frontend.nn import Module, Tensor, op, spec
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_nn_packing.py b/tests/python/relax/test_frontend_nn_packing.py
index 078cfdf66b..6342c17bd3 100644
--- a/tests/python/relax/test_frontend_nn_packing.py
+++ b/tests/python/relax/test_frontend_nn_packing.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _iter_binding_names(mod):
diff --git a/tests/python/relax/test_frontend_nn_subroutines.py b/tests/python/relax/test_frontend_nn_subroutines.py
index 0aa6a31532..11555869ca 100644
--- a/tests/python/relax/test_frontend_nn_subroutines.py
+++ b/tests/python/relax/test_frontend_nn_subroutines.py
@@ -21,7 +21,7 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_linear():
diff --git a/tests/python/relax/test_frontend_nn_tensor.py b/tests/python/relax/test_frontend_nn_tensor.py
index 5e2414fdb9..01e22cea37 100644
--- a/tests/python/relax/test_frontend_nn_tensor.py
+++ b/tests/python/relax/test_frontend_nn_tensor.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.nn import Module, Tensor, spec
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_tensor_from_numpy():
diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py
index edafb8124e..2956cdc0b2 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -44,8 +44,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
bg = np.random.MT19937(0)
diff --git a/tests/python/relax/test_frontend_stablehlo.py b/tests/python/relax/test_frontend_stablehlo.py
index 81f95ce0d7..adec4238ef 100644
--- a/tests/python/relax/test_frontend_stablehlo.py
+++ b/tests/python/relax/test_frontend_stablehlo.py
@@ -33,7 +33,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.stablehlo import from_stablehlo
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_inline_functions.py b/tests/python/relax/test_inline_functions.py
index c8c6b860e8..3deb8799ce 100644
--- a/tests/python/relax/test_inline_functions.py
+++ b/tests/python/relax/test_inline_functions.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_kill_after_last_use.py b/tests/python/relax/test_kill_after_last_use.py
index 0a37e7c3c3..8a21b44413 100644
--- a/tests/python/relax/test_kill_after_last_use.py
+++ b/tests/python/relax/test_kill_after_last_use.py
@@ -22,7 +22,7 @@ import tvm.relax
import tvm_tirx.testing, tvm
from tvm.relax.transform import KillAfterLastUse
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_lower_alloc_tensor.py b/tests/python/relax/test_lower_alloc_tensor.py
index 7cc6356906..928a317c4b 100644
--- a/tests/python/relax/test_lower_alloc_tensor.py
+++ b/tests/python/relax/test_lower_alloc_tensor.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LowerAllocTensor
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_basic():
diff --git a/tests/python/relax/test_meta_schedule_relax_integration.py b/tests/python/relax/test_meta_schedule_relax_integration.py
index 6cb286d1a0..8f8a35cb7e 100644
--- a/tests/python/relax/test_meta_schedule_relax_integration.py
+++ b/tests/python/relax/test_meta_schedule_relax_integration.py
@@ -30,7 +30,7 @@ from tvm_tirx.runtime import tensor as tvm_tensor
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_binary.py b/tests/python/relax/test_op_binary.py
index b1603a3fd4..b14f8a6d25 100644
--- a/tests/python/relax/test_op_binary.py
+++ b/tests/python/relax/test_op_binary.py
@@ -25,7 +25,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_ccl.py b/tests/python/relax/test_op_ccl.py
index f4533bbc74..3df7926737 100644
--- a/tests/python/relax/test_op_ccl.py
+++ b/tests/python/relax/test_op_ccl.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_create.py b/tests/python/relax/test_op_create.py
index 559cb9b64b..1f18ea2e54 100644
--- a/tests/python/relax/test_op_create.py
+++ b/tests/python/relax/test_op_create.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_datatype.py b/tests/python/relax/test_op_datatype.py
index bda623c0ee..bd1cb27811 100644
--- a/tests/python/relax/test_op_datatype.py
+++ b/tests/python/relax/test_op_datatype.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_grad.py b/tests/python/relax/test_op_grad.py
index d34a1bf99b..ab7dd792ba 100644
--- a/tests/python/relax/test_op_grad.py
+++ b/tests/python/relax/test_op_grad.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_image.py b/tests/python/relax/test_op_image.py
index 3cddf7d2ac..e90c16f02a 100644
--- a/tests/python/relax/test_op_image.py
+++ b/tests/python/relax/test_op_image.py
@@ -27,7 +27,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_index.py b/tests/python/relax/test_op_index.py
index 732b451686..0302ecbe5d 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_inspect.py b/tests/python/relax/test_op_inspect.py
index 0de689e1e9..18df1a09bd 100644
--- a/tests/python/relax/test_op_inspect.py
+++ b/tests/python/relax/test_op_inspect.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import Op
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# Parameterization for reading dtype of DLTensor. Chosen to have
diff --git a/tests/python/relax/test_op_linear_algebra.py b/tests/python/relax/test_op_linear_algebra.py
index 41ac265baa..66dd103926 100644
--- a/tests/python/relax/test_op_linear_algebra.py
+++ b/tests/python/relax/test_op_linear_algebra.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_manipulate.py b/tests/python/relax/test_op_manipulate.py
index 253719574f..d6748bad1b 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_misc.py b/tests/python/relax/test_op_misc.py
index 1684c065df..afb44600ef 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/tests/python/relax/test_op_misc.py
@@ -20,8 +20,8 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
index c2bcc276c0..1627282835 100644
--- a/tests/python/relax/test_op_nn.py
+++ b/tests/python/relax/test_op_nn.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_nn_convolution.py b/tests/python/relax/test_op_nn_convolution.py
index dbf76545dc..f6810b4b9b 100644
--- a/tests/python/relax/test_op_nn_convolution.py
+++ b/tests/python/relax/test_op_nn_convolution.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_conv1d_op_correctness():
diff --git a/tests/python/relax/test_op_nn_pooling.py b/tests/python/relax/test_op_nn_pooling.py
index 10de9d1f53..a69867ec78 100644
--- a/tests/python/relax/test_op_nn_pooling.py
+++ b/tests/python/relax/test_op_nn_pooling.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_qdq.py b/tests/python/relax/test_op_qdq.py
index 5d3026464d..bb40d1e653 100644
--- a/tests/python/relax/test_op_qdq.py
+++ b/tests/python/relax/test_op_qdq.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_sampling.py b/tests/python/relax/test_op_sampling.py
index e2076146c9..2600f73b5d 100644
--- a/tests/python/relax/test_op_sampling.py
+++ b/tests/python/relax/test_op_sampling.py
@@ -17,7 +17,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
diff --git a/tests/python/relax/test_op_search.py b/tests/python/relax/test_op_search.py
index 3e6703f509..2787e58d0c 100644
--- a/tests/python/relax/test_op_search.py
+++ b/tests/python/relax/test_op_search.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_set.py b/tests/python/relax/test_op_set.py
index 0660600432..6eb5864266 100644
--- a/tests/python/relax/test_op_set.py
+++ b/tests/python/relax/test_op_set.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_size.py b/tests/python/relax/test_op_size.py
index 222a6f07be..b58c16cb3e 100644
--- a/tests/python/relax/test_op_size.py
+++ b/tests/python/relax/test_op_size.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_size():
diff --git a/tests/python/relax/test_op_sort.py b/tests/python/relax/test_op_sort.py
index c4f62f27b8..c58dd9b1d6 100644
--- a/tests/python/relax/test_op_sort.py
+++ b/tests/python/relax/test_op_sort.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_statistical.py b/tests/python/relax/test_op_statistical.py
index 960ed1acda..b2b28d1ff7 100644
--- a/tests/python/relax/test_op_statistical.py
+++ b/tests/python/relax/test_op_statistical.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_take.py b/tests/python/relax/test_op_take.py
index b703015d27..6690d1ebb0 100644
--- a/tests/python/relax/test_op_take.py
+++ b/tests/python/relax/test_op_take.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
axis = tvm_tirx.testing.parameter(0, 1)
diff --git a/tests/python/relax/test_op_ternary.py b/tests/python/relax/test_op_ternary.py
index 8c94e68cd7..02a4991679 100644
--- a/tests/python/relax/test_op_ternary.py
+++ b/tests/python/relax/test_op_ternary.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_unary.py b/tests/python/relax/test_op_unary.py
index 25166f7d90..6bce7b83f9 100644
--- a/tests/python/relax/test_op_unary.py
+++ b/tests/python/relax/test_op_unary.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_view.py b/tests/python/relax/test_op_view.py
index eb86257472..421348beda 100644
--- a/tests/python/relax/test_op_view.py
+++ b/tests/python/relax/test_op_view.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_vision.py b/tests/python/relax/test_op_vision.py
index 65dc555478..a2ca7e67e0 100644
--- a/tests/python/relax/test_op_vision.py
+++ b/tests/python/relax/test_op_vision.py
@@ -29,7 +29,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py
index d3d4497059..adbe71785d 100644
--- a/tests/python/relax/test_optimize_layout_transform.py
+++ b/tests/python/relax/test_optimize_layout_transform.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform import DeadCodeElimination, FuseTIR, OptimizeLayoutTransform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py
index 9891a75e1b..f8cfe28a90 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_pytorch_integration.py b/tests/python/relax/test_pytorch_integration.py
index 0159865715..32d54036ae 100644
--- a/tests/python/relax/test_pytorch_integration.py
+++ b/tests/python/relax/test_pytorch_integration.py
@@ -32,8 +32,8 @@ import torch.nn.functional as F
import tvm, tvm_tirx
from tvm.relax import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n_matmul = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py
index 1be222ca9a..0733e9ec85 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -30,7 +30,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
diff --git a/tests/python/relax/test_relax_to_pyfunc_converter.py b/tests/python/relax/test_relax_to_pyfunc_converter.py
index 5555dd874e..8dd49fa9a8 100644
--- a/tests/python/relax/test_relax_to_pyfunc_converter.py
+++ b/tests/python/relax/test_relax_to_pyfunc_converter.py
@@ -30,8 +30,8 @@ import tvm, tvm_tirx
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.relax_to_pyfunc_converter import RelaxToPyFuncConverter
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n_symbolic_add = T.dynamic("n", "int64")
diff --git a/tests/python/relax/test_remove_redundant_reshape.py b/tests/python/relax/test_remove_redundant_reshape.py
index cdceca3689..f67fadb78e 100644
--- a/tests/python/relax/test_remove_redundant_reshape.py
+++ b/tests/python/relax/test_remove_redundant_reshape.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import DeadCodeElimination, RemoveRedundantReshape
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index 495042a56a..e7606cbccd 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_testing_nn.py b/tests/python/relax/test_testing_nn.py
index 512f8b6f55..42a81d6d59 100644
--- a/tests/python/relax/test_testing_nn.py
+++ b/tests/python/relax/test_testing_nn.py
@@ -20,7 +20,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.testing import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_tir_call_source_kernel.py b/tests/python/relax/test_tir_call_source_kernel.py
index b312a412e4..900f5f1d5e 100644
--- a/tests/python/relax/test_tir_call_source_kernel.py
+++ b/tests/python/relax/test_tir_call_source_kernel.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_training_append_loss.py b/tests/python/relax/test_training_append_loss.py
index 4048b378fc..03617986a5 100644
--- a/tests/python/relax/test_training_append_loss.py
+++ b/tests/python/relax/test_training_append_loss.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training import AppendLoss
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_loss.py b/tests/python/relax/test_training_loss.py
index 18680b6e3a..6c9c2af0c2 100644
--- a/tests/python/relax/test_training_loss.py
+++ b/tests/python/relax/test_training_loss.py
@@ -18,7 +18,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@I.ir_module
diff --git a/tests/python/relax/test_training_setup_trainer.py b/tests/python/relax/test_training_setup_trainer.py
index 8537657570..c655c581fb 100644
--- a/tests/python/relax/test_training_setup_trainer.py
+++ b/tests/python/relax/test_training_setup_trainer.py
@@ -25,7 +25,7 @@ from tvm.relax.training import SetupTrainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, MomentumSGD
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_trainer_numeric.py b/tests/python/relax/test_training_trainer_numeric.py
index 80f6b64113..72c72b4c0f 100644
--- a/tests/python/relax/test_training_trainer_numeric.py
+++ b/tests/python/relax/test_training_trainer_numeric.py
@@ -24,7 +24,7 @@ from tvm.relax.training import SetupTrainer, Trainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, Adam
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _get_backbone():
diff --git a/tests/python/relax/test_transform.py b/tests/python/relax/test_transform.py
index 9e7bf01625..c90559008e 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_allocate_workspace.py b/tests/python/relax/test_transform_allocate_workspace.py
index 20da882271..07ca383065 100644
--- a/tests/python/relax/test_transform_allocate_workspace.py
+++ b/tests/python/relax/test_transform_allocate_workspace.py
@@ -18,7 +18,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@I.ir_module
diff --git a/tests/python/relax/test_transform_alter_op_impl.py b/tests/python/relax/test_transform_alter_op_impl.py
index 784ac00c64..02139adf75 100644
--- a/tests/python/relax/test_transform_alter_op_impl.py
+++ b/tests/python/relax/test_transform_alter_op_impl.py
@@ -20,8 +20,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
kOperatorName = "operator_name"
diff --git a/tests/python/relax/test_transform_annotate_tir_op_pattern.py b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
index 163b71bc95..497a555ad9 100644
--- a/tests/python/relax/test_transform_annotate_tir_op_pattern.py
+++ b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
index 1fdcd642d2..625e187325 100644
--- a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
+++ b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
@@ -23,8 +23,8 @@ from tvm import relax
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_attach_global_symbol.py b/tests/python/relax/test_transform_attach_global_symbol.py
index 4b4ecbad2d..3542c98910 100644
--- a/tests/python/relax/test_transform_attach_global_symbol.py
+++ b/tests/python/relax/test_transform_attach_global_symbol.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bind_params.py b/tests/python/relax/test_transform_bind_params.py
index 6629e7b521..35af7e1218 100644
--- a/tests/python/relax/test_transform_bind_params.py
+++ b/tests/python/relax/test_transform_bind_params.py
@@ -23,8 +23,8 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
use_np_array = tvm_tirx.testing.parameter(False, True)
diff --git a/tests/python/relax/test_transform_bind_symbolic_vars.py b/tests/python/relax/test_transform_bind_symbolic_vars.py
index f2d7a42e9a..c4ba978eb9 100644
--- a/tests/python/relax/test_transform_bind_symbolic_vars.py
+++ b/tests/python/relax/test_transform_bind_symbolic_vars.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bundle_model_params.py b/tests/python/relax/test_transform_bundle_model_params.py
index 7ddbceefc7..b8e74bce66 100644
--- a/tests/python/relax/test_transform_bundle_model_params.py
+++ b/tests/python/relax/test_transform_bundle_model_params.py
@@ -27,7 +27,7 @@ import tvm_tirx.testing, tvm
import tvm.topi.testing
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py b/tests/python/relax/test_transform_canonicalize_bindings.py
index 8215e90052..fe4100f461 100644
--- a/tests/python/relax/test_transform_canonicalize_bindings.py
+++ b/tests/python/relax/test_transform_canonicalize_bindings.py
@@ -26,8 +26,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform.transform import CanonicalizeBindings
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_codegen_pass.py b/tests/python/relax/test_transform_codegen_pass.py
index 237b4088f5..dcfc56e323 100644
--- a/tests/python/relax/test_transform_codegen_pass.py
+++ b/tests/python/relax/test_transform_codegen_pass.py
@@ -30,8 +30,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.dpl import is_op, wildcard
from tvm.relax.testing import transform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import utils
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 36e673ab00..1aa8767bce 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_compute_prim_value.py b/tests/python/relax/test_transform_compute_prim_value.py
index a8d5660e9e..9ec51abd9a 100644
--- a/tests/python/relax/test_transform_compute_prim_value.py
+++ b/tests/python/relax/test_transform_compute_prim_value.py
@@ -18,8 +18,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_convert_dataflow.py b/tests/python/relax/test_transform_convert_dataflow.py
index 91e264f714..39e8f1587d 100644
--- a/tests/python/relax/test_transform_convert_dataflow.py
+++ b/tests/python/relax/test_transform_convert_dataflow.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
# functions that will not change
diff --git a/tests/python/relax/test_transform_cse.py b/tests/python/relax/test_transform_cse.py
index 9663268ee8..c776d11a60 100644
--- a/tests/python/relax/test_transform_cse.py
+++ b/tests/python/relax/test_transform_cse.py
@@ -22,7 +22,7 @@ import numpy as np
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import EliminateCommonSubexpr
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_dead_code_elimination.py b/tests/python/relax/test_transform_dead_code_elimination.py
index d79517c944..1a1a899051 100644
--- a/tests/python/relax/test_transform_dead_code_elimination.py
+++ b/tests/python/relax/test_transform_dead_code_elimination.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import DeadCodeElimination
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_decompose_ops.py b/tests/python/relax/test_transform_decompose_ops.py
index 32494d9a84..74bf4d82b3 100644
--- a/tests/python/relax/test_transform_decompose_ops.py
+++ b/tests/python/relax/test_transform_decompose_ops.py
@@ -25,7 +25,7 @@ from tvm_tirx import IRModule
from tvm import relax
from tvm.relax import Function
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_error_enrichment.py b/tests/python/relax/test_transform_error_enrichment.py
index dbf441eb41..f469386f9f 100644
--- a/tests/python/relax/test_transform_error_enrichment.py
+++ b/tests/python/relax/test_transform_error_enrichment.py
@@ -59,7 +59,7 @@ def test_pass_error_renders_underlined_tvmscript():
"Location (TVMScript):\n"
"Access path: <root>.body.blocks[0].bindings[0].value\n\n"
"from __future__ import annotations\n\n"
- "# from tvm_tirx.script import relax as R\n\n"
+ "# from tvm.relax import script as R\n\n"
"@R.function\n"
'def main(x: R.Tensor((3, 4), dtype="float32"), '
'y: R.Tensor((5, 6), dtype="float32")) -> R.Tensor((3, 6), dtype="float32"):\n'
diff --git a/tests/python/relax/test_transform_expand_matmul_of_sum.py b/tests/python/relax/test_transform_expand_matmul_of_sum.py
index f56421935c..23c0f74008 100644
--- a/tests/python/relax/test_transform_expand_matmul_of_sum.py
+++ b/tests/python/relax/test_transform_expand_matmul_of_sum.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_transform_expand_tuple_args.py b/tests/python/relax/test_transform_expand_tuple_args.py
index 441dfb8144..a8dde501e4 100644
--- a/tests/python/relax/test_transform_expand_tuple_args.py
+++ b/tests/python/relax/test_transform_expand_tuple_args.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
index 51a79a94bb..6556e96d0a 100644
--- a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
+++ b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
@@ -26,7 +26,7 @@ from tvm_tirx.ir.module import IRModule
from tvm.relax.expr_functor import PyExprVisitor, visitor
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_fold_constant.py b/tests/python/relax/test_transform_fold_constant.py
index 13f6e14089..7a5c60e822 100644
--- a/tests/python/relax/test_transform_fold_constant.py
+++ b/tests/python/relax/test_transform_fold_constant.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py
index 4b3cf4dfe2..7224458f3b 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
index 3d9b31e3da..e98e2681ba 100644
--- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
+++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
@@ -31,8 +31,8 @@ from tvm.relax.dpl.pattern import (
)
from tvm.relax.transform import PatternCheckContext
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py
index a562773624..83477501d3 100644
--- a/tests/python/relax/test_transform_fuse_tir.py
+++ b/tests/python/relax/test_transform_fuse_tir.py
@@ -22,8 +22,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_transpose_matmul.py b/tests/python/relax/test_transform_fuse_transpose_matmul.py
index 3924474906..031260d3aa 100644
--- a/tests/python/relax/test_transform_fuse_transpose_matmul.py
+++ b/tests/python/relax/test_transform_fuse_transpose_matmul.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_gradient.py b/tests/python/relax/test_transform_gradient.py
index 5d35b52c74..960ac3020f 100644
--- a/tests/python/relax/test_transform_gradient.py
+++ b/tests/python/relax/test_transform_gradient.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_gradient_te_register.py b/tests/python/relax/test_transform_gradient_te_register.py
index 29b429a368..c2966ba664 100644
--- a/tests/python/relax/test_transform_gradient_te_register.py
+++ b/tests/python/relax/test_transform_gradient_te_register.py
@@ -27,7 +27,7 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training.utils import register_te_gradient
from tvm.relax.transform import Gradient
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_inline_private_functions.py b/tests/python/relax/test_transform_inline_private_functions.py
index b59c40752c..78378d5eb4 100644
--- a/tests/python/relax/test_transform_inline_private_functions.py
+++ b/tests/python/relax/test_transform_inline_private_functions.py
@@ -23,7 +23,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
index 4db4f60f70..8370cc592a 100644
--- a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
+++ b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lambda_lift.py b/tests/python/relax/test_transform_lambda_lift.py
index a5ecfc352d..fb0cc2bc8f 100644
--- a/tests/python/relax/test_transform_lambda_lift.py
+++ b/tests/python/relax/test_transform_lambda_lift.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import transform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py b/tests/python/relax/test_transform_lazy_transform_params.py
index 3465230ef6..cde91a1c43 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -22,8 +22,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LazyTransformParams
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops.py b/tests/python/relax/test_transform_legalize_ops.py
index 73634c9be9..5303ba4eb5 100644
--- a/tests/python/relax/test_transform_legalize_ops.py
+++ b/tests/python/relax/test_transform_legalize_ops.py
@@ -26,8 +26,8 @@ from tvm import relax
from tvm.relax.transform import LegalizeOps
from tvm.relax.transform.legalize_ops.common import register_legalize
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_binary.py b/tests/python/relax/test_transform_legalize_ops_binary.py
index 484afd300a..bef48180e7 100644
--- a/tests/python/relax/test_transform_legalize_ops_binary.py
+++ b/tests/python/relax/test_transform_legalize_ops_binary.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Binary arithmetic #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_ccl.py b/tests/python/relax/test_transform_legalize_ops_ccl.py
index 12beaa1b80..45aeefcf25 100644
--- a/tests/python/relax/test_transform_legalize_ops_ccl.py
+++ b/tests/python/relax/test_transform_legalize_ops_ccl.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_create_datatype.py b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
index 9f0662df14..664b116200 100644
--- a/tests/python/relax/test_transform_legalize_ops_create_datatype.py
+++ b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
@@ -19,8 +19,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Creation #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_distributed.py b/tests/python/relax/test_transform_legalize_ops_distributed.py
index 87d670e166..213a2f3cd0 100644
--- a/tests/python/relax/test_transform_legalize_ops_distributed.py
+++ b/tests/python/relax/test_transform_legalize_ops_distributed.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_grad.py b/tests/python/relax/test_transform_legalize_ops_grad.py
index 7b3ce4ea3f..4ba0b113e7 100644
--- a/tests/python/relax/test_transform_legalize_ops_grad.py
+++ b/tests/python/relax/test_transform_legalize_ops_grad.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_image.py b/tests/python/relax/test_transform_legalize_ops_image.py
index daf0e85cc2..72982fa3b6 100644
--- a/tests/python/relax/test_transform_legalize_ops_image.py
+++ b/tests/python/relax/test_transform_legalize_ops_image.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
index ba26612000..81b279d5f7 100644
--- a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
+++ b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import Op
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Indexing #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index fb0888a5ca..639d9c6abd 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -23,8 +23,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Manipulation #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py b/tests/python/relax/test_transform_legalize_ops_nn.py
index 27b9d6428a..4528e13513 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Neural network #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_qdq.py b/tests/python/relax/test_transform_legalize_ops_qdq.py
index 0275aa7de8..3ac002ad33 100644
--- a/tests/python/relax/test_transform_legalize_ops_qdq.py
+++ b/tests/python/relax/test_transform_legalize_ops_qdq.py
@@ -18,8 +18,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_search_statistical.py b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
index f21d4a873c..7d59c9446e 100644
--- a/tests/python/relax/test_transform_legalize_ops_search_statistical.py
+++ b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Search #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py b/tests/python/relax/test_transform_legalize_ops_unary.py
index 8b5af413f9..ab88a99353 100644
--- a/tests/python/relax/test_transform_legalize_ops_unary.py
+++ b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm import topi
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lift_transform_params.py b/tests/python/relax/test_transform_lift_transform_params.py
index 4546b5d2e7..6f56df6511 100644
--- a/tests/python/relax/test_transform_lift_transform_params.py
+++ b/tests/python/relax/test_transform_lift_transform_params.py
@@ -29,8 +29,8 @@ pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py b/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
index d1db285ecf..8d5800e09e 100644
--- a/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
+++ b/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_merge_composite_functions.py b/tests/python/relax/test_transform_merge_composite_functions.py
index 05f9e322cb..7c4e9dd79d 100644
--- a/tests/python/relax/test_transform_merge_composite_functions.py
+++ b/tests/python/relax/test_transform_merge_composite_functions.py
@@ -19,8 +19,8 @@ import pytest
import tvm, tvm_tirx
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_meta_schedule_apply_database.py b/tests/python/relax/test_transform_meta_schedule_apply_database.py
index f6a397a8e9..8ad959ce21 100644
--- a/tests/python/relax/test_transform_meta_schedule_apply_database.py
+++ b/tests/python/relax/test_transform_meta_schedule_apply_database.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
target = tvm_tirx.target.Target({"kind": "llvm", "num-cores": 16})
diff --git a/tests/python/relax/test_transform_meta_schedule_tuning.py b/tests/python/relax/test_transform_meta_schedule_tuning.py
index decf585208..9592acd404 100644
--- a/tests/python/relax/test_transform_meta_schedule_tuning.py
+++ b/tests/python/relax/test_transform_meta_schedule_tuning.py
@@ -43,8 +43,8 @@ from tvm import relax
from tvm_tirx.ir import transform
from tvm_tirx.ir.module import IRModule
from tvm_tirx.ir.transform import PassContext
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
target = tvm_tirx.target.Target({"kind": "llvm", "num-cores": 16})
diff --git a/tests/python/relax/test_transform_normalize.py b/tests/python/relax/test_transform_normalize.py
index 46c9428a3d..0513b3f18a 100644
--- a/tests/python/relax/test_transform_normalize.py
+++ b/tests/python/relax/test_transform_normalize.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_normalize_global_var.py b/tests/python/relax/test_transform_normalize_global_var.py
index 29a53ce9cc..b8fa37e836 100644
--- a/tests/python/relax/test_transform_normalize_global_var.py
+++ b/tests/python/relax/test_transform_normalize_global_var.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_operator_specific_normalization.py b/tests/python/relax/test_transform_operator_specific_normalization.py
index e990d7eafb..55f598dc06 100644
--- a/tests/python/relax/test_transform_operator_specific_normalization.py
+++ b/tests/python/relax/test_transform_operator_specific_normalization.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm.relax.testing.transform
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_remove_unused_outputs.py b/tests/python/relax/test_transform_remove_unused_outputs.py
index 0fff54561e..19a700be16 100644
--- a/tests/python/relax/test_transform_remove_unused_outputs.py
+++ b/tests/python/relax/test_transform_remove_unused_outputs.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_remove_unused_parameters.py b/tests/python/relax/test_transform_remove_unused_parameters.py
index 16c2ffe5aa..870ce3f88d 100644
--- a/tests/python/relax/test_transform_remove_unused_parameters.py
+++ b/tests/python/relax/test_transform_remove_unused_parameters.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
index 23836513c0..257cd426ae 100644
--- a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
+++ b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_transform_reorder_take_after_matmul.py b/tests/python/relax/test_transform_reorder_take_after_matmul.py
index a7e1e6ba6a..3bc422cc82 100644
--- a/tests/python/relax/test_transform_reorder_take_after_matmul.py
+++ b/tests/python/relax/test_transform_reorder_take_after_matmul.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_rewrite_cuda_graph.py b/tests/python/relax/test_transform_rewrite_cuda_graph.py
index f209b41be1..51266413d3 100644
--- a/tests/python/relax/test_transform_rewrite_cuda_graph.py
+++ b/tests/python/relax/test_transform_rewrite_cuda_graph.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
index d8872a280e..27caeed1b5 100644
--- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
+++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
index 13e368d580..15899251e8 100644
--- a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
+++ b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
@@ -25,7 +25,7 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.module import IRModule
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_split_layout_rewrite_preproc.py b/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
index 955f1d8524..28bf5da290 100644
--- a/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
+++ b/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
@@ -19,8 +19,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_static_plan_block_memory.py b/tests/python/relax/test_transform_static_plan_block_memory.py
index 6a56e74fd7..e30c4cf56f 100644
--- a/tests/python/relax/test_transform_static_plan_block_memory.py
+++ b/tests/python/relax/test_transform_static_plan_block_memory.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_to_mixed_precision.py b/tests/python/relax/test_transform_to_mixed_precision.py
index 498f3f5741..b05dc19c68 100644
--- a/tests/python/relax/test_transform_to_mixed_precision.py
+++ b/tests/python/relax/test_transform_to_mixed_precision.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import ToMixedPrecision
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_topological_sort.py b/tests/python/relax/test_transform_topological_sort.py
index cf5c631f5d..fb85944fd6 100644
--- a/tests/python/relax/test_transform_topological_sort.py
+++ b/tests/python/relax/test_transform_topological_sort.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_depth_first_from_inputs():
diff --git a/tests/python/relax/test_transform_update_param_type.py b/tests/python/relax/test_transform_update_param_type.py
index fca9d83b5a..d49cc29719 100644
--- a/tests/python/relax/test_transform_update_param_type.py
+++ b/tests/python/relax/test_transform_update_param_type.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index 203f54e93a..e2c4616832 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -24,8 +24,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py
index 26bff991c0..13354415a4 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -36,8 +36,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.testing import nn
from tvm.relax.testing.vm import check_saved_func
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import cc, popen_pool, utils
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_vm_builtin.py b/tests/python/relax/test_vm_builtin.py
index c682e29be5..37489938ab 100644
--- a/tests/python/relax/test_vm_builtin.py
+++ b/tests/python/relax/test_vm_builtin.py
@@ -23,7 +23,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_multinomial_from_uniform():
diff --git a/tests/python/relax/test_vm_builtin_lower.py b/tests/python/relax/test_vm_builtin_lower.py
index f9c3a21295..dcb1a0bcb4 100644
--- a/tests/python/relax/test_vm_builtin_lower.py
+++ b/tests/python/relax/test_vm_builtin_lower.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.script, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_callback_function.py b/tests/python/relax/test_vm_callback_function.py
index b1621f9ba9..f1fecb574a 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
pytestmark = pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="llvm not enabled")
diff --git a/tests/python/relax/test_vm_codegen_only.py b/tests/python/relax/test_vm_codegen_only.py
index 866def31f1..950b672989 100644
--- a/tests/python/relax/test_vm_codegen_only.py
+++ b/tests/python/relax/test_vm_codegen_only.py
@@ -30,8 +30,8 @@ from tvm import relax
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm.relax.testing.vm import check_saved_func
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_cuda_graph.py b/tests/python/relax/test_vm_cuda_graph.py
index 65dfb73516..98379433a4 100644
--- a/tests/python/relax/test_vm_cuda_graph.py
+++ b/tests/python/relax/test_vm_cuda_graph.py
@@ -26,8 +26,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_vm_execbuilder.py b/tests/python/relax/test_vm_execbuilder.py
index 9f806f734f..af0a634290 100644
--- a/tests/python/relax/test_vm_execbuilder.py
+++ b/tests/python/relax/test_vm_execbuilder.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm, tvm_tirx
from tvm import relax
from tvm.relax.testing.vm import check_saved_func
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_vm_execute():
diff --git a/tests/python/relax/texture/test_texture_nd.py b/tests/python/relax/texture/test_texture_nd.py
index 0bce25b5c5..a1b0673215 100644
--- a/tests/python/relax/texture/test_texture_nd.py
+++ b/tests/python/relax/texture/test_texture_nd.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm.rpc import connect_tracker
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import ndk
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py b/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
index c4bc55fbcf..d6bade58b8 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
index f6d28a45cf..ed4b6f71ff 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
@@ -24,7 +24,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule
from tvm.s_tir.analysis import estimate_tir_flops
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
index 47e86ce13f..2addeb29d4 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir import StringImm, TensorRegion
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
identify_memcpy = tvm.s_tir.analysis._ffi_api._identify_memcpy
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
index 7c7721277d..8ce4130f6f 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
@@ -19,7 +19,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir.analysis import assert_pure_function, is_pure_function
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py b/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
index bf2e9bd3bb..5a4c7d8735 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
@@ -17,7 +17,7 @@
import pytest
import tvm, tvm_tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index 7a13518400..f8077628d2 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx.ir import TensorRegion, TensorRegionType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py
index 4410097d0b..bbcbffc9e5 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
index 3e476ae7df..668d4f4b3b 100644
--- a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
+++ b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
@@ -16,7 +16,7 @@
# under the License.
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/base/test_compilation_pipeline.py b/tests/python/s_tir/base/test_compilation_pipeline.py
index 3682dd3159..a0e0e4411c 100644
--- a/tests/python/s_tir/base/test_compilation_pipeline.py
+++ b/tests/python/s_tir/base/test_compilation_pipeline.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/base/test_s_tir_texture_scope.py b/tests/python/s_tir/base/test_s_tir_texture_scope.py
index 6425678045..40c911e596 100644
--- a/tests/python/s_tir/base/test_s_tir_texture_scope.py
+++ b/tests/python/s_tir/base/test_s_tir_texture_scope.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm_tirx.ir.module import IRModule
from tvm.s_tir.backend.adreno import pipeline as adreno_pipeline
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx._compile import _split_host_device_modules
from tvm_tirx.tirx.compilation_pipeline import finalize_device_passes
diff --git a/tests/python/s_tir/base/test_sblock_dependence_info.py b/tests/python/s_tir/base/test_sblock_dependence_info.py
index 78df1f4a08..81c3f9fe5a 100644
--- a/tests/python/s_tir/base/test_sblock_dependence_info.py
+++ b/tests/python/s_tir/base/test_sblock_dependence_info.py
@@ -28,7 +28,7 @@ from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
from tvm.s_tir import SBlockDependenceInfo
from tvm.s_tir.sblock_scope import DepKind
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Function
diff --git a/tests/python/s_tir/base/test_tir_te_extern_function.py b/tests/python/s_tir/base/test_tir_te_extern_function.py
index 57d4d77993..7cb1f41a80 100644
--- a/tests/python/s_tir/base/test_tir_te_extern_function.py
+++ b/tests/python/s_tir/base/test_tir_te_extern_function.py
@@ -24,7 +24,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import te
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# TODO(csullivan): Additional tests cases needed:
diff --git a/tests/python/s_tir/dlight/test_benchmark.py b/tests/python/s_tir/dlight/test_benchmark.py
index 4890e8b66f..5e563d0714 100644
--- a/tests/python/s_tir/dlight/test_benchmark.py
+++ b/tests/python/s_tir/dlight/test_benchmark.py
@@ -35,8 +35,8 @@ from tvm.s_tir.dlight.benchmark import (
)
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# The test function uses an undefined symbolic var in Relax.
diff --git a/tests/python/s_tir/dlight/test_cpu_gemv.py b/tests/python/s_tir/dlight/test_cpu_gemv.py
index 4b09008ef8..4503f5963a 100644
--- a/tests/python/s_tir/dlight/test_cpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_cpu_gemv.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_conv.py b/tests/python/s_tir/dlight/test_gpu_conv.py
index 2ffa34ce44..e84f3f6577 100644
--- a/tests/python/s_tir/dlight/test_gpu_conv.py
+++ b/tests/python/s_tir/dlight/test_gpu_conv.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_fallback.py b/tests/python/s_tir/dlight/test_gpu_fallback.py
index 5941cdf198..fe0d60fea4 100644
--- a/tests/python/s_tir/dlight/test_gpu_fallback.py
+++ b/tests/python/s_tir/dlight/test_gpu_fallback.py
@@ -23,7 +23,7 @@ from tvm import s_tir
from tvm_tirx.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py b/tests/python/s_tir/dlight/test_gpu_gemv.py
index 27db2044ab..3818c3b863 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -22,7 +22,7 @@ from __future__ import annotations
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_general_reduction.py b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
index bc0c23f826..7dd8fc42ce 100644
--- a/tests/python/s_tir/dlight/test_gpu_general_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index 183e49d853..7fdb62c704 100644
--- a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -22,7 +22,7 @@ from unittest import mock
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
from tvm.s_tir.dlight.gpu import low_batch_gemv
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul.py b/tests/python/s_tir/dlight/test_gpu_matmul.py
index 7fa571a494..a60cf7eeca 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index c0ed429e82..7c31b981f4 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_reduction.py b/tests/python/s_tir/dlight/test_gpu_reduction.py
index f280015656..09392212ee 100644
--- a/tests/python/s_tir/dlight/test_gpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_reduction.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_rmsnorm.py b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
index f389f9ad23..a5f8187057 100644
--- a/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
+++ b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
@@ -20,7 +20,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_transpose.py b/tests/python/s_tir/dlight/test_gpu_transpose.py
index 4ea67fd502..373384fd8f 100644
--- a/tests/python/s_tir/dlight/test_gpu_transpose.py
+++ b/tests/python/s_tir/dlight/test_gpu_transpose.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_primitives.py b/tests/python/s_tir/dlight/test_primitives.py
index 89728063ce..cb87931a53 100644
--- a/tests/python/s_tir/dlight/test_primitives.py
+++ b/tests/python/s_tir/dlight/test_primitives.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
index d1fe6614b2..1adad84751 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir.meta_schedule.arg_info import ArgInfo, TensorInfo
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
index 70fb6c8894..a6a31434a5 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
@@ -33,7 +33,7 @@ from tvm.s_tir.meta_schedule.builder import (
LocalBuilder,
PyBuilder,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
index 19f7d51098..d324cf2572 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
@@ -36,7 +36,7 @@ from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
requires_xgboost = pytest.mark.skipif(
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
index 2200f33e43..5ce1d9fa1a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
@@ -34,7 +34,7 @@ from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.database import TuningRecord, Workload
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
index 43268c1918..3c31bfea46 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
@@ -27,7 +27,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
N_FEATURES = 164
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
index 6256215174..5365a76cfd 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
index ed1217c528..003bb08134 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
index e2c371bf44..d11847c460 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
@@ -18,7 +18,7 @@
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
index cdfcaa20ad..2937fcd5ba 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
index 85a8a540a8..21a689fec1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
@@ -20,7 +20,7 @@ from functools import reduce
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
index f771b2488d..52f3f14906 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
@@ -18,7 +18,7 @@
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
index e7bfaf3350..bccfa40a11 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
@@ -33,7 +33,7 @@ from tvm.s_tir.meta_schedule.schedule_rule import PyScheduleRule
from tvm.s_tir.meta_schedule.space_generator import PostOrderApply
from tvm.s_tir.schedule import SBlockRV, Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
index c0805a106e..032348671a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
index b76daa3dc3..c8fdff1bcf 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
index 4350191143..88e6343c06 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
index 71362b1367..c8682407c2 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
index 4e8ba7dbae..bcefe183b1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm.s_tir.meta_schedule.postproc import RewriteParallelVectorizeUnroll
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
index 2787419b4f..e5696d822a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
index 4ee3612b85..fe1f5aee90 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.tensor_intrin import dot_product_common, x86
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
index d66098f084..7949c73878 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
index 113c3f8ad6..d65bdbd718 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
index 30de265812..3f2bae97c8 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
index 88e3827dcf..43316c34c1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
@@ -59,7 +59,7 @@ from tvm.s_tir.meta_schedule.utils import (
get_global_func_with_default_on_worker,
)
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.tirx import FloatImm
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
index 17d45e8bd6..2717ff6e86 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
index 8585caec88..da02f86654 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.schedule_rule import ApplyCustomRule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
index 219b4d4838..334f5ae704 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
index 403e150a2e..fd258b00d7 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
@@ -24,7 +24,7 @@ from tvm_tirx.ir.base import assert_structural_equal
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.space_generation import generate_design_space
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
index c9fe87f93f..2eac2a1ead 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
@@ -23,7 +23,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
index d488e49b9a..4e0f9c3fda 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
@@ -25,7 +25,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
generate_design_space,
print_sketches,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
index efb47fa766..72e67e126d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
@@ -27,7 +27,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
from tvm.s_tir.tensor_intrin.arm_cpu import DP4A_S8S8S32_INTRIN
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN as AVX512_INTRIN
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
index 00698287e3..766f8dc52d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
@@ -22,7 +22,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
index 4248eb33ef..623a549982 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
index ff0d22f977..282eea7b07 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
@@ -27,7 +27,7 @@ from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyMutator
from tvm.s_tir.schedule import Schedule, Trace
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
MATMUL_M = 32
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
index 63e8433038..f227a5efb1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
index 587a1fbf41..6de9d90fd4 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
index 3057559c88..d96701fee3 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
index 301f24ef25..709bb2e130 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
@@ -32,7 +32,7 @@ from tvm.s_tir.meta_schedule.space_generator import (
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
index c3097b329e..2cdf16e61b 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.runner.config import EvaluatorConfig
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
index 688f9a56fa..77eab1efc9 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
@@ -28,7 +28,7 @@ from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyBuilder, DummyRunner
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index a637fbae0c..0c7f8ed0de 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.codegen import llvm_lookup_intrinsic_id
from tvm.s_tir import Schedule
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.tirx import floordiv, floormod
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
index 53de2ba9ab..86c572aebc 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
@@ -25,7 +25,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.meta_schedule import TuneContext
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
index b2f21327ee..9940825376 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
@@ -30,7 +30,7 @@ from tvm.s_tir.meta_schedule.testing.custom_builder_runner import run_module_via
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm.s_tir.schedule import SBlockRV, Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index eb65388f75..18940a1ce7 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -34,7 +34,7 @@ from tvm.s_tir.schedule.analysis import (
suggest_index_map,
)
from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm.te import create_function
from tvm_tirx.tirx import IndexMap, Var, decl_tensor, floordiv, floormod
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py b/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
index 7e6fbab2f5..5b020ad86f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
@@ -23,7 +23,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
index b60150d72e..449ff60735 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule import DepKind
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_blockize.py b/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
index 00f89b494f..398baa0371 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py b/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
index 548b40326d..8ba7d03824 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
index 8a5188da65..c84d2e9680 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
index f94c2b57a7..2d0f7af255 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
index e74c2f0ce4..336c815430 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
index 4ac5183456..5515f928dd 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_error.py b/tests/python/s_tir/schedule/test_tir_schedule_error.py
index 115a46adfb..1e00c88cfb 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_error.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_error.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py b/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
index f2113aafe8..7fc2be231e 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
@@ -27,7 +27,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
index 890d28612f..b20c5f5e97 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
index b96119e176..96a5f40c8f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
index e40c333fe6..b98e66cb91 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_merge.py b/tests/python/s_tir/schedule/test_tir_schedule_merge.py
index b14f29c1ea..0729d78402 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_merge.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_merge.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py b/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
index c40f0ba642..28e3b836f3 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_partition.py b/tests/python/s_tir/schedule/test_tir_schedule_partition.py
index 23547a22f8..dfc117d188 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_partition.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_partition.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.expr import IntImm
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
index 28b09b2c3d..cc9e5e7089 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
@@ -43,7 +43,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
index c1da4be9f6..de645600cc 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
verify_trace_roundtrip,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
index a019bbf456..296cc6c3d1 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reorder.py b/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
index 50a5ec3d27..cb5ecdbaee 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
@@ -27,7 +27,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py b/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
index 8de2331719..ba55e0396b 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
index 6f40f26697..c756d13fc4 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
index d37b2ea507..ffbcd0c702 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_sampling.py b/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
index ddaa1a2c6b..12cff6aa2f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
@@ -24,7 +24,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py b/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
index 449304526f..c9f0fcc733 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
index a316c44566..588b79412a 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py b/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
index 76835bf875..a0a118b137 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.expr import IntImm
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state.py b/tests/python/s_tir/schedule/test_tir_schedule_state.py
index d820b4f786..e46b152288 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_state.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_state.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
@@ -95,7 +95,7 @@ def block_in_opaque_block(
def replace_ir_builder(deep_copy=False, realize=False):
new_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
s = tvm.s_tir.ScheduleState(new_func, debug_mask="all")
target = tvm.s_tir.SBlock(
@@ -124,11 +124,11 @@ def replace_ir_builder(deep_copy=False, realize=False):
def replace_ir_builder_module(deep_copy=False, realize=False):
new_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
other_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
mod = IRModule(functions={"main": new_func, "other": other_func})
s = tvm.s_tir.ScheduleState(mod, debug_mask="all")
@@ -158,7 +158,7 @@ def replace_ir_builder_module(deep_copy=False, realize=False):
def replace_ir_builder_with_opaque():
func = tvm_tirx.script.from_source(
block_in_opaque_block.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
s = tvm.s_tir.ScheduleState(func, debug_mask="all")
gc.collect()
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
index d296c81333..5bf2b2ea29 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule.state import CachedFlags
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
index dc6328f530..7041896fa0 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
index 05f7f433c6..db254cf127 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
@@ -38,7 +38,7 @@ from tvm.s_tir.tensor_intrin.arm_cpu import (
from tvm.s_tir.tensor_intrin.hexagon import VDMPY_i16i16i32_INTRIN, VRMPY_u8u8i32_INTRIN
from tvm.s_tir.tensor_intrin.rocm import AMDGPU_SDOT4_INTRIN
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_trace.py b/tests/python/s_tir/schedule/test_tir_schedule_trace.py
index cd2a0683f6..d3e6b25fe4 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_trace.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_trace.py
@@ -27,7 +27,7 @@ from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule import Instruction, InstructionKind, LoopRV, SBlockRV, Trace
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform.py b/tests/python/s_tir/schedule/test_tir_schedule_transform.py
index a3d338ca68..d50cb42d7e 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform.py
@@ -18,7 +18,7 @@ import tvm, tvm_tirx
from tvm.s_tir import Schedule
from tvm.s_tir.schedule.transform import tile_with_tensor_intrin
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
index a4b24a4b9d..331e037a88 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
verify_trace_roundtrip,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_utilities.py b/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
index e6e3c52691..0921c276a6 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py b/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
index 7dc932f6d5..53e823e736 100644
--- a/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
+++ b/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index b3fe568b85..eb61d0d404 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import from_source
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -474,8 +474,8 @@ def test_for_thread_binding():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -761,8 +761,8 @@ def test_assert_stmt_roundtrip(error_type, message_parts):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
@@ -956,7 +956,7 @@ def func():
j{annotation} = i // 4
T.evaluate(j)
""",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
binding = func.body[0].body.seq[0]
if mutable:
@@ -1076,8 +1076,8 @@ def test_roundtrip_basic_usage(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_blocks.py b/tests/python/s_tir/script/test_s_tir_script_blocks.py
index 843d662f7f..8a99e5c18e 100644
--- a/tests/python/s_tir/script/test_s_tir_script_blocks.py
+++ b/tests/python/s_tir/script/test_s_tir_script_blocks.py
@@ -25,7 +25,7 @@ import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# Expected completed types use the configured fixed index dtype, not shape dtype.
@@ -305,14 +305,14 @@ def expected_recursive_tensorslice_indices(
def test_complete_tensor_indices():
new_func = tvm_tirx.script.from_source(
func_with_tensorslice_indices.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
new_func, expected_tensorslice_indices.with_attr("global_symbol", "main")
)
new_func = tvm_tirx.script.from_source(
func_with_recursive_tensorslice_indices.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
new_func, expected_recursive_tensorslice_indices.with_attr("global_symbol", "main")
@@ -389,7 +389,7 @@ def expect_alloc_tensor_func(
def test_complete_alloc_tensor():
rt_func = tvm_tirx.script.from_source(
alloc_tensor_func.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
rt_func, expect_alloc_tensor_func.with_attr("global_symbol", "main")
@@ -443,11 +443,11 @@ def test_alloc_zero_dim_tensor_round_trip():
func = alloc_zero_dim_tensor
func_with_block = alloc_zero_dim_tensor_block
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
rt_func_with_block = tvm_tirx.script.from_source(
func_with_block.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
rt_mod = tvm.compile(rt_func, "llvm")
rt_mod_with_block = tvm.compile(rt_func_with_block, "llvm")
@@ -641,8 +641,8 @@ def test_matmul_original():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -663,8 +663,8 @@ def test_element_wise():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -687,8 +687,8 @@ def test_predicate():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -726,8 +726,8 @@ def test_match_tensor_region():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -776,8 +776,8 @@ def test_block_elements():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -817,8 +817,8 @@ def test_opaque_block():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -1025,8 +1025,8 @@ def test_roundtrip_blocks(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
index 2d56718dac..9ed2c42e93 100644
--- a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
+++ b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
@@ -25,7 +25,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import from_source
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -105,7 +105,7 @@ def gemm_dyn_shape(
def test_dynamic_shape_gemm():
gemm_dyn_shape_roundtrip = from_source(
gemm_dyn_shape.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
assert_structural_equal_ignore_global_symbol(gemm_dyn_shape, gemm_dyn_shape_roundtrip)
@@ -273,8 +273,8 @@ def test_roundtrip_dynamic_shape(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_error_handling.py b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
index b02c66f341..91e1dc83c1 100644
--- a/tests/python/s_tir/script/test_s_tir_script_error_handling.py
+++ b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
@@ -30,7 +30,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import from_source
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
index 59599cfa93..b2387b5d2e 100644
--- a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
+++ b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
@@ -25,7 +25,7 @@ from tvm import s_tir
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.ir.location import SourceLoc, SourceName
from tvm.s_tir.script.ir_builder import function as build_function
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_meta_programming.py b/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
index fef353325f..596b18e8b8 100644
--- a/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
+++ b/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -506,7 +506,7 @@ def matmul(
def test_multi_element_array_in_outmost_namespace():
func = matmul
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
diff --git a/tests/python/s_tir/script/test_s_tir_script_namespace.py b/tests/python/s_tir/script/test_s_tir_script_namespace.py
index dcc05b300e..4dbd103c38 100644
--- a/tests/python/s_tir/script/test_s_tir_script_namespace.py
+++ b/tests/python/s_tir/script/test_s_tir_script_namespace.py
@@ -26,13 +26,13 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires Python 3.12")
def test_shared_operations_and_aliases():
- from tvm_tirx.script import s_tir as S
+ from tvm.s_tir import script as S
n = S.dynamic("n")
@@ -53,7 +53,7 @@ def test_shared_operations_and_aliases():
shared,
tvm_tirx.script.from_source(
shared.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
)
@@ -82,7 +82,7 @@ def test_mixed_module_roundtrip():
tvm_tirx.ir.assert_structural_equal(
Mixed,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ script, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
)
direct_script = Mixed["direct"].script()
@@ -106,7 +106,7 @@ def test_tirx_construction_roundtrip_and_execution_are_independent(monkeypatch):
restored = tvm_tirx.script.from_source(
Direct.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
tvm_tirx.ir.assert_structural_equal(Direct, restored)
assert "s_tir" not in restored["main"].attrs
@@ -136,7 +136,7 @@ def test_legacy_mode_is_not_a_function_option(namespace, option):
tvm_tirx.script.from_source(
f"@{namespace}.function({option}=True, check_well_formed=False)\n"
"def main():\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -169,7 +169,7 @@ def test_tirx_rejects_s_tir_blocks():
tvm_tirx.script.from_source(
"@T.function(check_well_formed=False)\n"
"def main():\n with Ts.sblock('bad'):\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -179,7 +179,7 @@ def test_tirx_cannot_change_dialect_with_attribute(value):
tvm_tirx.script.from_source(
"@T.function(check_well_formed=False)\n"
f"def main():\n T.func_attr({{'s_tir': {value}}})\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -211,7 +211,7 @@ def test_s_tir_options_and_helpers():
scheduled,
tvm_tirx.script.from_source(
scheduled.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
)
@@ -223,7 +223,7 @@ def test_s_tir_options_and_helpers():
"tvm_tirx.tirx.script",
"tvm_tirx.script.parser",
"tvm.s_tir.script.ir_builder.frame",
- "tvm_tirx.script.ir_builder.s_tir.frame",
+ "tvm.s_tir.script.ir_builder",
],
)
def test_import_order(first):
@@ -233,7 +233,7 @@ def test_import_order(first):
"-c",
f"import {first}\n"
"import tvm.s_tir.script as direct\n"
- "from tvm_tirx.script import s_tir as Ts, tirx as T\n"
+ "from tvm.s_tir import script as Ts\nfrom tvm_tirx.script import tirx as T\n"
"from tvm_tirx.script.parser import s_tir as parser\n"
"assert Ts is direct is parser\n"
"assert Ts.Tensor is T.Tensor\n"
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py
index c85a6675cb..92bc63d305 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -30,7 +30,7 @@ from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
from tvm.s_tir.script.ir_builder import function as build_function
from tvm_tirx.script import ir as I
from tvm_tirx.script import ir_builder as IB
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as TB
@@ -55,7 +55,7 @@ def test_function_symbolic_tensor_param_roundtrip():
assert source.index("def main[n](") < source.index("T.evaluate(n)")
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
func,
)
@@ -76,7 +76,7 @@ def test_function_compound_tensor_shape_first_use_roundtrip():
assert source.index("def main[n](") < source.index("T.evaluate(n)")
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
func,
)
@@ -109,7 +109,7 @@ def test_function_symbolic_alloc_tensor_roundtrip():
tvm_tirx.script.from_source(
source,
check_well_formed=False,
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
func,
)
@@ -132,7 +132,7 @@ def test_function():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -159,7 +159,7 @@ def test_function_tensor_data_use():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -188,7 +188,7 @@ def test_function_tensor_data_argument_is_scope_hint():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -286,7 +286,7 @@ def test_bind():
_assert_print(
obj,
"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -322,7 +322,7 @@ def test_remap():
v4, v5 = Ts.axis.remap("RS", [i4, i5])
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -370,7 +370,7 @@ def test_root_block():
TB.evaluate(0)
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -403,7 +403,7 @@ def test_private_function():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -424,7 +424,7 @@ def test_function_different_symbol():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -442,7 +442,7 @@ def test_return_statement():
return TB.int32(5)
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -481,7 +481,7 @@ def test_custom_float_types(dtype):
TB.evaluate(getattr(TB, dtype)(0.0))
expected_output = f"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -523,7 +523,7 @@ def test_predicated_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -570,7 +570,7 @@ def test_predicated_tensor_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -607,7 +607,7 @@ def test_predicated_scalable_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -633,7 +633,7 @@ def test_masked_load_prevents_scalar_allocation_init_fusion():
assert "A[0] = T.masked_load" in source
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
main,
)
@@ -649,7 +649,7 @@ def test_vload_with_explicit_scalable_data_type():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -672,7 +672,7 @@ def test_vectorize_llvm_pure_intrin():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -698,7 +698,7 @@ def test_func_with_loop_jumps():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -1272,8 +1272,8 @@ def test_same_name_var():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -1458,8 +1458,8 @@ def test_div_mod():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func, True)
@@ -1698,8 +1698,8 @@ def test_void_ptr_vs_handle():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
),
)
@@ -1718,8 +1718,8 @@ def test_void_ptr_vs_handle():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
),
)
@@ -2072,7 +2072,7 @@ def test_annotation_multi_access_paths():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2100,7 +2100,7 @@ def test_annotate_from_multi_obj():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2133,7 +2133,7 @@ def test_disable_concise_scoping_when_scope_annotated():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2163,7 +2163,7 @@ def test_ir_module():
_assert_module_representations(
mod,
"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -2223,8 +2223,8 @@ def test_roundtrip_lowered(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2258,8 +2258,8 @@ def test_roundtrip_expressions(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2291,8 +2291,8 @@ def test_roundtrip_tensors(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2316,8 +2316,8 @@ def test_roundtrip_metadata(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/s_tir/script/test_s_tir_script_source_locations.py
index 37f1f49c0f..3de5b52938 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/s_tir/script/test_s_tir_script_source_locations.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/test_arith_domain_touched.py b/tests/python/s_tir/test_arith_domain_touched.py
index 82571ea8b0..3a7394b77f 100644
--- a/tests/python/s_tir/test_arith_domain_touched.py
+++ b/tests/python/s_tir/test_arith_domain_touched.py
@@ -20,7 +20,7 @@ import pytest
import tvm_ffi
import tvm, tvm_tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int32")
diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py
index 26105985e6..06210d940a 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir import Var
from tvm.s_tir import SBlock
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py b/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
index fb17cb8a84..c1f6741770 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py b/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
index 48c11d681b..0290952d7a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
index 33bd10cac5..834c03b0ce 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
index 03e51bec07..305934a423 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
index a93cf084d1..ab0a0d558b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.transform import DefaultGPUSchedule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
index 9a61db7898..66f8ed9d0d 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
@@ -18,7 +18,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
index 0d2c6b76c7..be44ac40d1 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
@@ -25,7 +25,7 @@ import tvm.s_tir.testing
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -42,7 +42,7 @@ def test_cp_async_raw_dtype_round_trips():
T.s_tir.cp_async_raw(B.data, i * 16, A.data, i * 16, 16, ty="float16")
reparsed = tvm_tirx.script.from_source(
- f.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ f.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(f, reparsed)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index bff50ff2c0..4e0fb7c267 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import te
from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index d8d483822c..2f43c21c7a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -20,7 +20,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py b/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
index f052f00110..2937a64721 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
index 1651d8c5da..c8c0e1ff61 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
@@ -22,7 +22,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
index 1f8f04d210..eaece8d2cd 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
@@ -23,7 +23,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
index bd47d236e7..2ba8464acc 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
@@ -17,7 +17,7 @@
# ruff: noqa: F401, F821
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-self-argument
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
index f44d283866..b236799b05 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm.s_tir
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index 546cc3b104..504c98eb64 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -19,7 +19,7 @@ from __future__ import annotations
import tvm, tvm_tirx
import tvm.s_tir
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
index e46c2f6ae2..bf9bfacd67 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
index 1b8a390319..94930c2cf3 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py b/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
index 3a39f77407..3940f2a509 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
index 18981bfe84..725dbf929b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
index 63af21b391..37283e1f55 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm.topi.math import cast
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
index 957bc52a33..5b58bc1b1e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
index bd6d7e08d8..6ec0bf9a68 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.tensor_intrin.hexagon import VRMPY_u8u8i32_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
index dc2357864a..c366a9a8fe 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py b/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
index a0ebd683b9..dfbaee591e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
@@ -19,7 +19,7 @@ import sys
import tvm, tvm_tirx
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
index cce7cd3826..7d43632576 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index cc205c6864..bbf76290c3 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/script/test_constructor_contracts.py b/tests/python/script/test_constructor_contracts.py
index 6cb90862d4..a99d765641 100644
--- a/tests/python/script/test_constructor_contracts.py
+++ b/tests/python/script/test_constructor_contracts.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import ir, tirx
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/target/test_arm_target.py b/tests/python/target/test_arm_target.py
index b0c3a6834b..acefce9a7f 100644
--- a/tests/python/target/test_arm_target.py
+++ b/tests/python/target/test_arm_target.py
@@ -25,7 +25,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import codegen
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/target/test_target_target.py b/tests/python/target/test_target_target.py
index 1ad749d77a..d9d9210f19 100644
--- a/tests/python/target/test_target_target.py
+++ b/tests/python/target/test_target_target.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/te/test_te_create_function.py b/tests/python/te/test_te_create_function.py
index a4c90b9d25..8811415e40 100644
--- a/tests/python/te/test_te_create_function.py
+++ b/tests/python/te/test_te_create_function.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te, topi
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -471,7 +471,7 @@ def test_tensor_attr():
)
func = te.create_function([A, B, C])
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/tirx/codegen/conftest.py
similarity index 60%
copy from tests/python/s_tir/script/test_s_tir_script_source_locations.py
copy to tests/python/tirx/codegen/conftest.py
index 37f1f49c0f..dabfb09775 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/tirx/codegen/conftest.py
@@ -14,28 +14,20 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""S-TIR script source locations."""
+"""Hardware requirements for TIRx codegen tests."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
-from tvm_tirx.script import tirx as T
+from pathlib import Path
+from tvm_tirx.testing._pytest import (
+ _release_cuda_cache_between_tests as _release_cuda_cache_between_tests,
+)
+from tvm_tirx.testing._pytest import (
+ _source_codegen_without_cuda as _source_codegen_without_cuda,
+)
+from tvm_tirx.testing._pytest import (
+ codegen_collection_modifyitems,
+)
-def test_nesting_parsing():
- class dummy:
- pass
- for i in range(1):
-
- @tvm_tirx.script.ir_module
- class Module:
- @Ts.function
- def impl(
- A: T.Tensor((12, 196, 64), "float32"),
- ) -> None:
- T.evaluate(0)
-
-
-if __name__ == "__main__":
- tvm_tirx.testing.main()
+def pytest_collection_modifyitems(config, items):
+ codegen_collection_modifyitems(config, items, Path(__file__).resolve().parent)
diff --git a/tests/python/tirx/codegen/test_tirx_codegen_fp8.py b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
index 5838144065..1b5bfb3a49 100644
--- a/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
+++ b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
@@ -19,7 +19,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/tirx/conftest.py
similarity index 60%
copy from tests/python/s_tir/script/test_s_tir_script_source_locations.py
copy to tests/python/tirx/conftest.py
index 37f1f49c0f..8e18f0bd93 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/tirx/conftest.py
@@ -14,28 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""S-TIR script source locations."""
+"""Shared CUDA device and exact-architecture handling for TIRx tests."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
-from tvm_tirx.script import tirx as T
-
-
-def test_nesting_parsing():
- class dummy:
- pass
-
- for i in range(1):
-
- @tvm_tirx.script.ir_module
- class Module:
- @Ts.function
- def impl(
- A: T.Tensor((12, 196, 64), "float32"),
- ) -> None:
- T.evaluate(0)
-
-
-if __name__ == "__main__":
- tvm_tirx.testing.main()
+from tvm_tirx.testing import _pytest as _pytest
+from tvm_tirx.testing._pytest import pytest_configure as pytest_configure
+from tvm_tirx.testing._pytest import pytest_runtest_setup as pytest_runtest_setup
diff --git a/tirx/tests/python/tirx/conftest.py b/tirx/python/tvm_tirx/testing/_pytest.py
similarity index 79%
copy from tirx/tests/python/tirx/conftest.py
copy to tirx/python/tvm_tirx/testing/_pytest.py
index fa896876f5..89bd9497ff 100644
--- a/tirx/tests/python/tirx/conftest.py
+++ b/tirx/python/tvm_tirx/testing/_pytest.py
@@ -16,8 +16,10 @@
# under the License.
"""Shared CUDA device and exact-architecture handling for TIRx tests."""
+import gc
import os
import re
+from pathlib import Path
import pytest
@@ -110,3 +112,35 @@ def pytest_runtest_setup(item):
f"requires CUDA architecture {required} on {device_source}; "
f"actual architecture is {actual or 'unavailable'}"
)
+
+
+@pytest.fixture(autouse=True)
+def _source_codegen_without_cuda(monkeypatch):
+ # CPU-only source checks must not depend on a CUDA compiler being installed.
+ if not env.has_cuda():
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
+
+@pytest.fixture(autouse=True)
+def _release_cuda_cache_between_tests():
+ yield
+ gc.collect()
+ try:
+ import torch
+ except ImportError:
+ return
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+
+def codegen_collection_modifyitems(config, items, suite_root):
+ if env.has_cuda_compute(10):
+ return
+ suite_root = Path(suite_root).resolve()
+ skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
+ for item in items:
+ if (
+ Path(item.path).resolve().is_relative_to(suite_root)
+ and item.get_closest_marker("gpu") is not None
+ ):
+ item.add_marker(skip)
diff --git a/tirx/tests/python/runtime/test_runtime_memory_scope.py b/tirx/tests/python/runtime/test_runtime_memory_scope.py
index 6f0e0c7e55..23602b279d 100644
--- a/tirx/tests/python/runtime/test_runtime_memory_scope.py
+++ b/tirx/tests/python/runtime/test_runtime_memory_scope.py
@@ -126,7 +126,7 @@ def test_opencl_memory_scope_registration_load_order(load_order):
)
if not opencl_path.is_file():
pytest.skip("OpenCL runtime sidecar is not built")
- compiler = tvm_tirx.base._LOADED_LIBS.get("tvm_compiler")
+ compiler = tvm_tirx.base._LOADED_LIBS.get("tvm_tirx")
if compiler is None and load_order != "runtime_only":
pytest.skip("Compiler library is not available")
env = {
diff --git a/tirx/tests/python/tirx/codegen/conftest.py b/tirx/tests/python/tirx/codegen/conftest.py
index 769f579282..dabfb09775 100644
--- a/tirx/tests/python/tirx/codegen/conftest.py
+++ b/tirx/tests/python/tirx/codegen/conftest.py
@@ -16,41 +16,18 @@
# under the License.
"""Hardware requirements for TIRx codegen tests."""
-import gc
from pathlib import Path
-import pytest
-
-from tvm_tirx.testing import env
-
-
-@pytest.fixture(autouse=True)
-def _source_codegen_without_cuda(monkeypatch):
- # CPU-only source checks must not depend on a CUDA compiler being installed.
- if not env.has_cuda():
- monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
-
-
-@pytest.fixture(autouse=True)
-def _release_cuda_cache_between_tests():
- yield
- gc.collect()
- try:
- import torch
- except ImportError:
- return
- if torch.cuda.is_available():
- torch.cuda.empty_cache()
+from tvm_tirx.testing._pytest import (
+ _release_cuda_cache_between_tests as _release_cuda_cache_between_tests,
+)
+from tvm_tirx.testing._pytest import (
+ _source_codegen_without_cuda as _source_codegen_without_cuda,
+)
+from tvm_tirx.testing._pytest import (
+ codegen_collection_modifyitems,
+)
def pytest_collection_modifyitems(config, items):
- if env.has_cuda_compute(10):
- return
- suite_root = Path(__file__).resolve().parent
- skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
- for item in items:
- if (
- Path(item.path).resolve().is_relative_to(suite_root)
- and item.get_closest_marker("gpu") is not None
- ):
- item.add_marker(skip)
+ codegen_collection_modifyitems(config, items, Path(__file__).resolve().parent)
diff --git a/tirx/tests/python/tirx/conftest.py b/tirx/tests/python/tirx/conftest.py
index fa896876f5..8e18f0bd93 100644
--- a/tirx/tests/python/tirx/conftest.py
+++ b/tirx/tests/python/tirx/conftest.py
@@ -16,97 +16,6 @@
# under the License.
"""Shared CUDA device and exact-architecture handling for TIRx tests."""
-import os
-import re
-
-import pytest
-
-from tvm_tirx.testing import env
-
-_CUDA_ARCH_PATTERN = re.compile(r"sm_[1-9][0-9]*(?:a|f)?\Z")
-_XDIST_CUDA_DEVICE = None
-
-
-def _visible_cuda_archs():
- try:
- import torch
- except ImportError:
- return ()
-
- if not torch.cuda.is_available():
- return ()
- return tuple(env.cuda_arch(device_id) for device_id in range(torch.cuda.device_count()))
-
-
-def _require_homogeneous_cuda_archs():
- arches = _visible_cuda_archs()
- if not arches:
- return None
- if any(arch is None for arch in arches):
- raise pytest.UsageError(f"could not determine every visible CUDA architecture: {arches}")
- if len(set(arches)) != 1:
- profile = ", ".join(f"cuda:{device_id}={arch}" for device_id, arch in enumerate(arches))
- raise pytest.UsageError(
- "mixed-architecture CUDA pools are unsupported by the TIRx test suite: " + profile
- )
- return arches[0]
-
-
-def _set_cuda_device_for_xdist_worker():
- global _XDIST_CUDA_DEVICE
-
- try:
- import torch
- except ImportError:
- return None
-
- if not torch.cuda.is_available():
- return None
- if _XDIST_CUDA_DEVICE is None:
- worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
- worker_index = int(worker[2:]) if worker.startswith("gw") and worker[2:].isdigit() else 0
- _XDIST_CUDA_DEVICE = worker_index % torch.cuda.device_count()
- torch.cuda.set_device(_XDIST_CUDA_DEVICE)
- return _XDIST_CUDA_DEVICE
-
-
-def pytest_configure(config):
- del config
- _require_homogeneous_cuda_archs()
- _set_cuda_device_for_xdist_worker()
-
-
-def pytest_runtest_setup(item):
- current_device = _set_cuda_device_for_xdist_worker()
- marker = item.get_closest_marker("cuda_arch")
- if marker is None:
- return
-
- unknown_kwargs = set(marker.kwargs) - {"device"}
- if unknown_kwargs:
- raise pytest.UsageError(
- f"cuda_arch marker has unsupported keyword(s): {sorted(unknown_kwargs)}"
- )
- arches = tuple(marker.args)
- if not arches or any(
- not isinstance(arch, str) or _CUDA_ARCH_PATTERN.fullmatch(arch) is None for arch in arches
- ):
- raise pytest.UsageError(
- "cuda_arch marker requires one or more canonical architectures such as sm_100a"
- )
-
- device_source = marker.kwargs.get("device", "cuda0")
- if device_source == "cuda0":
- device_id = 0
- elif device_source == "current":
- device_id = current_device if current_device is not None else 0
- else:
- raise pytest.UsageError("cuda_arch marker device must be 'cuda0' or 'current'")
-
- actual = env.cuda_arch(device_id)
- if actual not in arches:
- required = ", ".join(arches)
- pytest.skip(
- f"requires CUDA architecture {required} on {device_source}; "
- f"actual architecture is {actual or 'unavailable'}"
- )
+from tvm_tirx.testing import _pytest as _pytest
+from tvm_tirx.testing._pytest import pytest_configure as pytest_configure
+from tvm_tirx.testing._pytest import pytest_runtest_setup as pytest_runtest_setup
diff --git a/tirx/tests/python/tirx/script/test_tirx_script_tracing.py b/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
index cde7d9e796..bea151d4ac 100644
--- a/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
+++ b/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
@@ -301,7 +301,7 @@ def test_shared_tensor_shapes_are_checked_by_compiled_host():
I.arg_("x", T.Tensor((n,), T.float32))
I.arg_("y", T.Tensor((n,), T.float32))
I.return_(n)
- compiled = tvm_tirx.compile(builder.get(), target="llvm", tir_pipeline="tirx")
+ compiled = tvm_tirx.tirx.compile(builder.get(), target="llvm", pipeline="tirx")
assert compiled(torch.zeros(3), torch.zeros(3)) == 3
for other, error, message in (
(torch.zeros(4), ValueError, "expected to match"),
diff --git a/tirx/tests/python/tirx/test_cuda_arch_marker.py b/tirx/tests/python/tirx/test_cuda_arch_marker.py
index 261657b424..c571924b07 100644
--- a/tirx/tests/python/tirx/test_cuda_arch_marker.py
+++ b/tirx/tests/python/tirx/test_cuda_arch_marker.py
@@ -27,7 +27,7 @@ def _tirx_conftest(pytestconfig):
for plugin in pytestconfig.pluginmanager.get_plugins():
plugin_path = getattr(plugin, "__file__", None)
if plugin_path is not None and Path(plugin_path).resolve() == expected:
- return plugin
+ return plugin._pytest
raise AssertionError(f"TIRx conftest plugin {expected} is not loaded")
diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py
index df65a07b10..6c7ad26852 100644
--- a/tirx/tests/python/tirx/test_cuda_backend_config.py
+++ b/tirx/tests/python/tirx/test_cuda_backend_config.py
@@ -656,7 +656,7 @@ def main(A: T.Tensor((32,), "float32")):
target = TARGET
if target == "target_object":
target = tvm_tirx.target.Target("cuda", host="llvm")
-module = tvm_tirx.compile(main, target=target)
+module = tvm_tirx.tirx.compile(main, target=target)
metadata = module.mod.imports[0].inspect_source("backend_config")
from tvm_tirx.backend.config import parse_backend_config
assert parse_backend_config(metadata)["cuda"]["arch"] == "sm_100a"
diff --git a/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
index 9fd84cd39d..f8c8b8f168 100644
--- a/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
+++ b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
@@ -168,7 +168,7 @@ def test_lower_floordiv_and_floormod(dtype, kind, variant):
[tvm_tirx.tirx.transform.StmtSimplify(), tvm_tirx.tirx.transform.LowerIntrin()]
)(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- compiled = tvm_tirx.compile(After, "llvm")
+ compiled = tvm_tirx.tirx.compile(After, "llvm")
data = itertools.product(range(-10, 10), [y for y in range(-10, 10) if y])
for x, y in data:
if (variant == "pow2" and y != 8) or (variant == "offset" and y != 5):
@@ -226,7 +226,7 @@ def test_lower_floordiv_overflow_checks(dtype, clamp, data):
[tvm_tirx.tirx.transform.StmtSimplify(), tvm_tirx.tirx.transform.LowerIntrin()]
)(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- compiled = tvm_tirx.compile(After, "llvm")
+ compiled = tvm_tirx.tirx.compile(After, "llvm")
for value in data:
expected = min(max(value, -10), 32758) if clamp else value
assert compiled(value) == expected // 3