__init__.py
 1  """TVM End-to-End Optimization Module.
 2  
 3  This module provides utilities for optimizing deep learning models using Apache TVM,
 4  from PyTorch model import to TensorIR manual scheduling.
 5  """
 6  
 7  import contextlib
 8  
 9  from .baseline import (
10      convert_to_relay,
11      load_pytorch_model,
12      run_baseline_inference,
13      run_pytorch_baseline,
14  )
15  
16  with contextlib.suppress(ImportError):
17      from .auto_scheduler import optimize_with_auto_scheduler, run_auto_tuned_inference
18  
19  with contextlib.suppress(ImportError):
20      from .tensorir_schedule import create_gemm_schedule, run_tensorir_gemm
21  
22  __version__ = "1.0.0"
23  __all__ = [
24      "convert_to_relay",
25      "create_gemm_schedule",
26      "load_pytorch_model",
27      "optimize_with_auto_scheduler",
28      "run_auto_tuned_inference",
29      "run_baseline_inference",
30      "run_pytorch_baseline",
31      "run_tensorir_gemm",
32  ]