# isort: skip_file
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, wrong-import-position
"""The Relax IR namespace containing the IR, type, operator, builder, vm, etc."""

from tvm.runtime import vm
from tvm.runtime.vm import VirtualMachine, VMInstrumentReturnKind
from tvm.ir import Call

# Expr
from .expr import (
    Expr,
    Span,
    GlobalVar,
    Var,
    DataflowVar,
    Binding,
    MatchCast,
    VarBinding,
    BindingBlock,
    DataflowBlock,
    SeqExpr,
    ShapeExpr,
    Tuple,
    TupleGetItem,
    Function,
    ExternFunc,
    If,
    Constant,
    DataTypeImm,
    StringImm,
    prim_value,
)

from .expr import const, extern, get_shape_of

# Type
from .type import (
    Type,
    AnyType,
    ObjectType,
    ShapeType,
    TensorType,
    TupleType,
    FuncType,
    PackedFuncType,
)

# VM
from .exec_builder import ExecBuilder

# Operator
from .op.base import (
    call_tir,
    call_tir_inplace,
    call_pure_packed,
    call_dps_packed,
    call_tir_with_grad,
)

# BlockBuilder
from .block_builder import BlockBuilder

# ExprFunctor
from .expr_functor import ExprFunctor, PyExprVisitor, PyExprMutator

# pipeline
from .pipeline import get_default_pipeline
from .pipeline import get_pipeline
from .pipeline import register_pipeline

# utils
from .utils import convert_to_expr

# BasePyModule
from .base_py_module import BasePyModule

# Import submodules in the last to avoid dependency
from . import exec_builder
from . import expr
from . import ty
from . import type
from . import analysis
from . import transform
from . import block_builder
from . import op
from . import backend
from . import training
from . import distributed
from . import frontend
from . import utils

# VM
from .vm_build import build, VMExecutable

from .binding_rewrite import DataflowBlockRewrite

import tvm.script

tvm.script.register_dialect("relax", "tvm.relax.script")