import sys
_module = sys.modules[__name__]
del sys
check_index_rst = _module
doc_link_checker = _module
quantize_to_ncnn = _module
test_java_demo = _module
test_onnx2ncnn = _module
ascend = _module
coreml = _module
ncnn = _module
onnxruntime = _module
openvino = _module
pplnn = _module
rknn = _module
sdk = _module
snpe = _module
tensorrt = _module
torchscript = _module
tvm = _module
onnx_config = _module
torchscript_config = _module
classification_dynamic = _module
classification_ncnn_dynamic = _module
classification_ncnn_static = _module
classification_onnxruntime_dynamic = _module
classification_onnxruntime_static = _module
classification_sdk_dynamic = _module
classification_snpe_static = _module
classification_static = _module
classification_torchscript = _module
base_dynamic = _module
base_static = _module
base_torchscript = _module
detection_onnxruntime_dynamic = _module
detection_onnxruntime_static = _module
detection_sdk_dynamic = _module
detection_torchscript = _module
yolov3_partition_onnxruntime_static = _module
segmentation_dynamic = _module
segmentation_onnxruntime_dynamic = _module
segmentation_sdk_dynamic = _module
segmentation_static = _module
segmentation_torchscript = _module
demo_rewrite = _module
image_classification = _module
image_restorer = _module
image_segmentation = _module
object_detection = _module
ocr = _module
pipeline = _module
pose_detection = _module
rotated_object_detection = _module
video_recognition = _module
conf = _module
mmdeploy = _module
apis = _module
calibration = _module
core = _module
pipeline_manager = _module
extract_model = _module
inference = _module
onnx = _module
export = _module
optimizer = _module
partition = _module
passes = _module
optimize_onnx = _module
utils = _module
pytorch2onnx = _module
pytorch2torchscript = _module
torch_jit = _module
trace = _module
calibration = _module
visualize = _module
backend = _module
backend_manager = _module
onnx2ascend = _module
wrapper = _module
base = _module
backend_wrapper_registry = _module
base_wrapper = _module
ops = _module
torchscript2coreml = _module
wrapper = _module
init_plugins = _module
onnx2ncnn = _module
quant = _module
wrapper = _module
wrapper = _module
onnx2openvino = _module
wrapper = _module
onnx2pplnn = _module
wrapper = _module
onnx2rknn = _module
wrapper = _module
export_info = _module
tracer = _module
onnx2dlc = _module
wrapper = _module
calib_utils = _module
onnx2tensorrt = _module
torch_allocator = _module
utils = _module
wrapper = _module
backend_manager = _module
wrapper = _module
onnx2tvm = _module
quantize = _module
tuner = _module
wrapper = _module
codebase = _module
backend_model = _module
mmcodebase = _module
task = _module
mmaction = _module
deploy = _module
mmaction = _module
video_recognition = _module
video_recognition_model = _module
models = _module
recognizers = _module
base = _module
recognizer2d = _module
recognizer3d = _module
mmcls = _module
classification = _module
classification_model = _module
mmclassification = _module
backbones = _module
shufflenet_v2 = _module
vision_transformer = _module
classifiers = _module
heads = _module
cls_head = _module
multi_label_head = _module
necks = _module
gap = _module
attention = _module
mmdet = _module
anchor = _module
bbox = _module
delta_xywh_bbox_coder = _module
distance_point_bbox_coder = _module
tblr_bbox_coder = _module
transforms = _module
detection_output = _module
prior_box = _module
point_generator = _module
post_processing = _module
bbox_nms = _module
mmdetection = _module
model_partition_cfg = _module
object_detection = _module
object_detection_model = _module
utils = _module
backbones = _module
dense_heads = _module
base_dense_head = _module
detr_head = _module
fovea_head = _module
gfl_head = _module
reppoints_head = _module
rpn_head = _module
ssd_head = _module
yolo_head = _module
yolox_head = _module
detectors = _module
base = _module
rpn = _module
single_stage = _module
two_stage = _module
necks = _module
roi_heads = _module
bbox_head = _module
cascade_roi_head = _module
fcn_mask_head = _module
single_level_roi_extractor = _module
standard_roi_head = _module
test_mixins = _module
transformer = _module
mmdet3d = _module
anchor_3d_generator = _module
centerpoint_bbox_coders = _module
fcos3d_bbox_coder = _module
utils = _module
box3d_nms = _module
mmdetection3d = _module
monocular_detection = _module
monocular_detection_model = _module
voxel_detection = _module
voxel_detection_model = _module
anchor3d_head = _module
anchor_free_mono3d_head = _module
centerpoint = _module
fcos_mono3d = _module
fcos_mono3d_head = _module
mvx_two_stage = _module
pillar_encode = _module
pillar_scatter = _module
single_stage_mono3d = _module
voxelnet = _module
mmedit = _module
mmediting = _module
super_resolution = _module
super_resolution_model = _module
mmocr = _module
mmocr = _module
text_detection = _module
text_detection_model = _module
text_recognition = _module
text_recognition_model = _module
fpn_cat = _module
single_stage_text_detector = _module
base = _module
crnn_decoder = _module
encode_decode_recognizer = _module
lstm_layer = _module
sar = _module
sar_decoder = _module
sar_encoder = _module
mmpose = _module
mmpose = _module
pose_detection = _module
pose_detection_model = _module
top_down = _module
deeppose_regression_head = _module
topdown_heatmap_multi_stage_head = _module
topdown_heatmap_simple_head = _module
vipnas_heatmap_simple_head = _module
mmrotate = _module
delta_midpointoffset_rbbox_coder = _module
delta_xywha_rbbox_coder = _module
gliding_vertex_coder = _module
transforms = _module
bbox_nms = _module
mmrotate = _module
rotated_detection = _module
rotated_detection_model = _module
oriented_rpn_head = _module
rotated_anchor_head = _module
rotated_rpn_head = _module
gv_bbox_head = _module
gv_ratio_roi_head = _module
oriented_standard_roi_head = _module
roi_extractors = _module
roi_trans_roi_head = _module
rotated_bbox_head = _module
single_stage_rotated_detector = _module
mmseg = _module
mmsegmentation = _module
segmentation = _module
segmentation_model = _module
decode_heads = _module
aspp_head = _module
ema_head = _module
point_head = _module
segmentors = _module
base = _module
encoder_decoder = _module
up_conv_block = _module
optimizers = _module
extractor = _module
function_marker = _module
optimize = _module
rewriters = _module
function_rewriter = _module
module_rewriter = _module
rewriter_manager = _module
rewriter_utils = _module
symbolic_rewriter = _module
mmcv = _module
cnn = _module
context_block = _module
hsigmoid = _module
hswish = _module
transformer = _module
deform_conv = _module
modulated_deform_conv = _module
nms = _module
nms_rotated = _module
point_sample = _module
roi_align = _module
roi_align_rotated = _module
pytorch = _module
functions = _module
adaptive_pool = _module
atan2 = _module
chunk = _module
clip = _module
expand = _module
flatten = _module
getattribute = _module
group_norm = _module
interpolate = _module
linear = _module
masked_fill = _module
mod = _module
multi_head_attention_forward = _module
normalize = _module
repeat = _module
size = _module
tensor_getitem = _module
tensor_setitem = _module
topk = _module
triu = _module
gelu = _module
grid_sampler = _module
hardsigmoid = _module
instance_norm = _module
layer_norm = _module
linear = _module
lstm = _module
pad = _module
roll = _module
squeeze = _module
config_utils = _module
constants = _module
dataset = _module
device = _module
env = _module
logging = _module
test = _module
timer = _module
utils = _module
version = _module
inference_pb2 = _module
inference_pb2_grpc = _module
setup = _module
srgan = _module
test_calibration = _module
test_extract = _module
test_onnx2ascend = _module
test_onnx2ncnn = _module
test_onnx2openvino = _module
test_onnx2rknn = _module
test_onnx2tensorrt = _module
test_onnx2tvm = _module
test_onnx_passes = _module
test_torch2onnx = _module
test_torch2torchscript = _module
test_wrapper = _module
mmrazor_model = _module
model = _module
test_classification = _module
test_classification_model = _module
test_mmcls_models = _module
test_mmdet_core = _module
test_mmdet_models = _module
test_mmdet_utils = _module
test_object_detection = _module
test_object_detection_model = _module
model_cfg = _module
monodet_model_cfg = _module
test_mmdet3d_models = _module
test_monocular_detection = _module
test_monocular_detection_model = _module
test_voxel_detection = _module
test_voxel_detection_model = _module
test_mmedit_models = _module
test_super_resolution = _module
test_super_resolution_model = _module
crnn = _module
dbnet = _module
test_mmocr_models = _module
test_text_detection = _module
test_text_detection_models = _module
test_text_recognition = _module
test_text_recognition_models = _module
test_mmpose_models = _module
test_pose_detection = _module
test_pose_detection_model = _module
test_mmrotate_core = _module
test_mmrotate_models = _module
test_rotated_detection = _module
test_rotated_detection_model = _module
test_mmseg_models = _module
test_segmentation = _module
test_segmentation_model = _module
package = _module
module = _module
test_function_rewriter = _module
test_mark = _module
test_module_rewriter = _module
test_rewriter_utils = _module
test_symbolic_register = _module
test_mmcv_cnn = _module
test_mmcv_ops = _module
test_ops = _module
test_ops = _module
utils = _module
test_pytorch_functions = _module
test_pytorch_ops = _module
test_dataset = _module
test_timer = _module
test_util = _module
check_env = _module
deploy = _module
extract_transform = _module
extract = _module
onnx2ncnn_quant_table = _module
mmdeploy_builder = _module
mmdeploy_python = _module
profiler = _module
quant_image_dataset = _module
regression_test = _module
build_ubuntu_x64_ncnn = _module
build_ubuntu_x64_ort = _module
build_ubuntu_x64_pplnn = _module
build_ubuntu_x64_torchscript = _module
build_ubuntu_x64_tvm = _module
ubuntu_utils = _module
sdk_analyze = _module
torch2onnx = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import torch


from torchvision.models import resnet18


from typing import Optional


from typing import Union


import inspect


import logging


from functools import wraps


from typing import Any


from typing import Callable


from typing import Dict


from typing import List


from typing import Sequence


import numpy as np


from copy import deepcopy


from functools import partial


from typing import Tuple


from torch.utils.data import DataLoader


from typing import NamedTuple


from abc import ABCMeta


from abc import abstractmethod


import abc


import time


from random import randint


import re


from torch.utils.data import Dataset


from torch import Tensor


from torch.onnx import symbolic_helper


import torch.nn.functional as F


from torch.nn import functional as F


from torch.autograd import Function


import torch.nn as nn


import warnings


from torch import nn


import copy


from typing import MutableSequence


from torch.onnx.symbolic_helper import parse_args


from torch.onnx import symbolic_helper as sym_help


from torch.nn.modules.utils import _pair


from torch.types import Number


import math


from typing import Iterable


from torch.onnx.symbolic_helper import _get_tensor_dim_size


from torch.onnx.symbolic_helper import _get_tensor_rank


from torch.onnx.symbolic_helper import _unimplemented


from torch.onnx.symbolic_helper import _unsqueeze_helper


import torch.onnx.symbolic_helper as sym_help


from torch.onnx.symbolic_opset9 import unused


from torch.onnx.symbolic_helper import _slice_helper


import random


import string


from torch.utils.data.dataset import Dataset


import torch.multiprocessing as mp


from torch.multiprocessing import Process


from torch.multiprocessing import set_start_method


import pandas as pd


from torch.hub import download_url_to_file


class BaseWrapper(torch.nn.Module, metaclass=ABCMeta):
    """Abstract base class for backend wrappers.

    Args:
        output_names (Sequence[str]): Names to model outputs in order, which is
        useful when converting the output dict to a ordered list or converting
        the output ordered list to a key-value dict.
    """

    def __init__(self, output_names: Sequence[str]):
        super().__init__()
        self._output_names = output_names

    @staticmethod
    def get_backend_file_count() ->int:
        """Return the count of backend file(s)

        Each backend has its own requirement on backend files (e.g., TensorRT
        requires 1 .engine file and ncnn requires 2 files (.param, .bin)). This
        interface allow developers to get the count of these required files.

        Returns:
            int: The count of required backend file(s).
        """
        return 1

    @abstractmethod
    def forward(self, inputs: Dict[str, torch.Tensor]) ->Dict[str, torch.Tensor]:
        """Run forward inference.

        Args:
            inputs (Dict[str, torch.Tensor]): Key-value pairs of model inputs.

        Returns:
            Dict[str, torch.Tensor]: Key-value pairs of model outputs.
        """
        pass

    @property
    def output_names(self):
        """Return the output names."""
        return self._output_names

    @output_names.setter
    def output_names(self, value):
        """Set the output names."""
        self._output_names = value

    def output_to_list(self, output_dict: Dict[str, torch.Tensor]) ->List[torch.Tensor]:
        """Convert the output dict of forward() to a tensor list.

        Args:
            output_dict (Dict[str, torch.Tensor]): Key-value pairs of model
                outputs.

        Returns:
            List[torch.Tensor]: An output value list whose order is determined
                by the ouput_names list.
        """
        outputs = [output_dict[name] for name in self._output_names]
        return outputs


logger_initialized = {}


def get_logger(name: str, log_file: Optional[str]=None, log_level: int=logging.INFO, file_mode: str='w'):
    """Initialize and get a logger by name.

    If the logger has not been initialized, this method will initialize the
    logger by adding one or two handlers, otherwise the initialized logger will
    be directly returned. During initialization, a StreamHandler will always be
    added. If `log_file` is specified, a FileHandler will also be added.
    Args:
        name (str): Logger name.
        log_file (str | None): The log filename. If specified, a FileHandler
            will be added to the logger.
        log_level (int): The logger level.
        file_mode (str): The file mode used in opening log file.
            Defaults to 'w'.
    Returns:
        logging.Logger: The expected logger.
    """
    logger = logging.getLogger(name)
    if name in logger_initialized:
        return logger
    for logger_name in logger_initialized:
        if name.startswith(logger_name):
            return logger
    for handler in logger.root.handlers:
        if type(handler) is logging.StreamHandler:
            handler.setLevel(logging.ERROR)
    stream_handler = logging.StreamHandler()
    handlers = [stream_handler]
    if log_file is not None:
        file_handler = logging.FileHandler(log_file, file_mode)
        handlers.append(file_handler)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    for handler in handlers:
        handler.setFormatter(formatter)
        handler.setLevel(log_level)
        logger.addHandler(handler)
    logger.setLevel(log_level)
    logger_initialized[name] = True
    return logger


def get_root_logger(log_file=None, log_level=logging.INFO) ->logging.Logger:
    """Get root logger.

    Args:
        log_file (str, optional): File path of log. Defaults to None.
        log_level (int, optional): The level of logger.
            Defaults to logging.INFO.
    Returns:
        logging.Logger: The obtained logger
    """
    logger = get_logger(name='mmdeploy', log_file=log_file, log_level=log_level)
    return logger


def get_file_path(prefix, candidates) ->str:
    """Search for file in candidates.

    Args:
        prefix (str): Prefix of the paths.
        candidates (str): Candidate paths
    Returns:
        str: file path or '' if not found
    """
    for candidate in candidates:
        wildcard = os.path.abspath(os.path.join(prefix, candidate))
        paths = glob.glob(wildcard)
        if paths:
            lib_path = paths[0]
            return lib_path
    return ''


def get_ops_path() ->str:
    """Get ncnn custom ops library path.

    Returns:
        str: The library path of ncnn custom ops.
    """
    candidates = ['../../lib/libmmdeploy_ncnn_ops.so', '../../lib/mmdeploy_ncnn_ops.dll']
    return get_file_path(os.path.dirname(__file__), candidates)


def parse_cuda_device_id(device: str) ->int:
    """Parse cuda device index from a string.

    Args:
        device (str): The typical style of string specifying cuda device,
            e.g.: 'cuda:0'.

    Returns:
        int: The parsed device id, defaults to `0`.
    """
    match_result = re.match('([^:]+)(:[0-9]+)?$', device)
    assert match_result is not None, f'Can not parse device {device}.'
    assert match_result.group(1).lower() == 'cuda', 'Not cuda device.'
    device_id = 0 if match_result.lastindex == 1 else int(match_result.group(2)[1:])
    return device_id


def parse_device_id(device: str) ->Optional[int]:
    """Parse device index from a string.

    Args:
        device (str): The typical style of string specifying device,
            e.g.: 'cuda:0', 'cpu'.

    Returns:
        Optional[int]: The return value depends on the type of device.
            If device is 'cuda': cuda device index, defaults to `0`.
            If device is 'cpu': `-1`.
            Otherwise, `None` will be returned.
    """
    if device == 'cpu':
        return -1
    if 'cuda' in device:
        return parse_cuda_device_id(device)
    return None


class SleepingPolicy(abc.ABC):

    @abc.abstractmethod
    def sleep(self, try_i: int):
        """How long to sleep in milliseconds.

        :param try_i: the number of retry (starting from zero)
        """
        assert try_i >= 0


class ExponentialBackoff(SleepingPolicy):

    def __init__(self, *, init_backoff_ms: int, max_backoff_ms: int, multiplier: int):
        self.init_backoff = randint(0, init_backoff_ms)
        self.max_backoff = max_backoff_ms
        self.multiplier = multiplier

    def sleep(self, try_i: int):
        sleep_range = min(self.init_backoff * self.multiplier ** try_i, self.max_backoff)
        sleep_ms = randint(0, sleep_range)
        logger = get_root_logger()
        logger.debug(f'Sleeping for {sleep_ms}')
        time.sleep(sleep_ms / 1000)


def get_env_key() ->str:
    """Return environment key str.

    Returns:
        str: The string to find SNPE service URI
    """
    return '__MMDEPLOY_SNPE_URI'


def load_tensorrt_plugin() ->bool:
    """Load TensorRT plugins library.

    Returns:
        bool: True if TensorRT plugin library is successfully loaded.
    """
    lib_path = get_ops_path()
    success = False
    logger = get_root_logger()
    if os.path.exists(lib_path):
        ctypes.CDLL(lib_path)
        logger.info(f'Successfully loaded tensorrt plugins from {lib_path}')
        success = True
    else:
        logger.warning(f'Could not load the library of tensorrt plugins.Because the file does not exist: {lib_path}')
    return success


class Error(Exception):
    """Acl Exception."""
    pass


def _check(code: int, msg: str):
    """check the error code.

    Args:
        code (int): The error code.
        msg (str): Error message.
    """
    if code != 0:
        raise Error(msg, code)


_is_torch_npu_available = False


class Context:
    ref_count = 0
    owned_acl = False

    def __init__(self):
        if not _is_torch_npu_available:
            self._active = True
            if Context.ref_count == 0:
                ret = acl.init()
                if ret == 0:
                    Context.owned_acl = True
                elif ret == 100002:
                    pass
                else:
                    _check(ret, 'acl.init')
            Context.ref_count += 1
        else:
            self._active = False

    def __del__(self):
        self.destroy()

    def destroy(self):
        if not self._active:
            return
        Context.ref_count -= 1
        if Context.ref_count == 0 and Context.owned_acl:
            ret = acl.finalize()
            if ret == 0:
                Context.owned_acl = False
            elif ret == 100037:
                pass
            else:
                _check(ret, 'acl.finalize')
        self._active = False


class TraceFunc:
    """Trace Transform."""

    def __init__(self):
        self.module_dict = dict()

    def register_module(self, name):
        if name in self.module_dict:
            raise KeyError(f'{name} is already registered')

        def _register(func):
            self.module_dict[name] = func
            return func
        return _register

    def get(self, name):
        return self.module_dict[name]


_TRANSFORM_WRAPPER = TraceFunc()


@_TRANSFORM_WRAPPER.register_module(name='Resize')
def resize(context: Context, args: Dict):
    context.transforms.append({'type': 'Resize'})
    return True


class ONNXNMSop(torch.autograd.Function):
    """Create onnx::NonMaxSuppression op.

    NMS in mmcv only supports one class with no batch info. This class assists
    in exporting NMS of ONNX's definition.
    """

    @staticmethod
    def forward(ctx, boxes: Tensor, scores: Tensor, max_output_boxes_per_class: int, iou_threshold: float, score_threshold: float) ->Tensor:
        """Get NMS output indices.

        Args:
            ctx (Context): The context with meta information.
            boxes (Tensor): The bounding boxes of shape [N, num_boxes, 4].
            scores (Tensor): The detection scores of shape
                [N, num_boxes, num_classes].
            max_output_boxes_per_class (int): Maximum number of output
                boxes per class of nms.
            iou_threshold (float): IOU threshold of nms.
            score_threshold (float): score threshold of nms.

        Returns:
            Tensor: Selected indices of boxes. 2-D tensor of shape
            (num_selected_indices, 3) with each row of
            [batch_index, class_index, box_index].
        """
        batch_size, num_class, _ = scores.shape
        score_threshold = float(score_threshold)
        iou_threshold = float(iou_threshold)
        indices = []
        for batch_id in range(batch_size):
            for cls_id in range(num_class):
                _boxes = boxes[batch_id, ...]
                _scores = scores[batch_id, cls_id, ...].contiguous()
                _, box_inds = nms(_boxes, _scores, iou_threshold, offset=0, score_threshold=score_threshold, max_num=max_output_boxes_per_class)
                batch_inds = torch.zeros_like(box_inds) + batch_id
                cls_inds = torch.zeros_like(box_inds) + cls_id
                indices.append(torch.stack([batch_inds, cls_inds, box_inds], dim=-1))
        indices = torch.cat(indices)
        return indices

    @staticmethod
    def symbolic(g, boxes: Tensor, scores: Tensor, max_output_boxes_per_class: int, iou_threshold: float, score_threshold: float):
        """Symbolic function for onnx::NonMaxSuppression.

        Args:
            g (Graph): The traced onnx graph.
            boxes (Tensor): The bounding boxes of shape [N, num_boxes, 4].
            scores (Tensor): The detection scores of shape
                [N, num_boxes, num_classes].
            max_output_boxes_per_class (int): Maximum number of output
                boxes per class of nms.
            iou_threshold (float): IOU threshold of nms.
            score_threshold (float): score threshold of nms.

        Returns:
            NonMaxSuppression op for onnx.
        """
        if not sym_help._is_value(max_output_boxes_per_class):
            max_output_boxes_per_class = g.op('Constant', value_t=torch.tensor(max_output_boxes_per_class, dtype=torch.long))
        if not sym_help._is_value(iou_threshold):
            iou_threshold = g.op('Constant', value_t=torch.tensor([iou_threshold], dtype=torch.float))
        if not sym_help._is_value(score_threshold):
            score_threshold = g.op('Constant', value_t=torch.tensor([score_threshold], dtype=torch.float))
        return g.op('NonMaxSuppression', boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold)


def select_nms_index(scores, bboxes, nms_index, keep_top_k, dir_scores=None, attr_scores=None):
    """Transform NMSBEV output.
    Args:
        scores (Tensor): The detection scores of shape
            [N, num_classes, num_boxes].
        boxes (Tensor): The bounding boxes of shape [N, num_boxes, 6].
        nms_index (Tensor): NMS output of bounding boxes indexing.
        keep_top_k (int): Number of top K boxes to keep after nms.
            Defaults to -1.
        dir_scores (Tensor): The direction scores of shape
            [N, num_boxes]. Defaults to None.
        attr_scores (Tensor): The attribute scores of shape
            [N, num_boxes]. Defaults to None.
    Returns:
        tuple[Tensor, Tensor, Tensor]: (bbox, scores, labels), `dets` of
            shape [N, num_det, 9] and `scores` of shape [N, num_det] and
            `labels` of shape [N, num_det]
    """
    batch_inds, cls_inds = nms_index[:, 0], nms_index[:, 1]
    box_inds = nms_index[:, 2]
    scores = scores[batch_inds, cls_inds, box_inds].unsqueeze(0)
    bboxes = bboxes[batch_inds, box_inds, ...].unsqueeze(0)
    labels = cls_inds.unsqueeze(0)
    if dir_scores is not None:
        dir_scores = dir_scores[batch_inds, box_inds].unsqueeze(0)
    if attr_scores is not None:
        attr_scores = attr_scores[batch_inds, box_inds].unsqueeze(0)
    is_use_topk = keep_top_k > 0 and (torch.onnx.is_in_onnx_export() or keep_top_k < scores.shape[1])
    if is_use_topk:
        scores, topk_inds = scores.topk(keep_top_k, dim=1)
    else:
        scores, topk_inds = scores.sort(dim=1, descending=True)
    topk_inds = topk_inds.squeeze(0)
    bboxes = bboxes[:, topk_inds, :]
    labels = labels[:, topk_inds]
    results = bboxes, scores, labels
    if dir_scores is not None:
        dir_scores = dir_scores[:, topk_inds]
        results = results + (dir_scores,)
    if attr_scores is not None:
        attr_scores = attr_scores[:, topk_inds]
        results = results + (attr_scores,)
    return results


version = __version__


def _multiclass_nms(boxes: Tensor, scores: Tensor, max_output_boxes_per_class: int=1000, iou_threshold: float=0.5, score_threshold: float=0.05, pre_top_k: int=-1, keep_top_k: int=-1):
    """Create a dummy onnx::NonMaxSuppression op while exporting to ONNX.

    This function helps exporting to onnx with batch and multiclass NMS op. It
    only supports class-agnostic detection results. That is, the scores is of
    shape (N, num_bboxes, num_classes) and the boxes is of shape (N, num_boxes,
    4).
    """
    if version.parse(torch.__version__) < version.parse('1.13.0'):
        max_output_boxes_per_class = torch.LongTensor([max_output_boxes_per_class])
    iou_threshold = torch.tensor([iou_threshold], dtype=torch.float32)
    score_threshold = torch.tensor([score_threshold], dtype=torch.float32)
    batch_size = scores.shape[0]
    if pre_top_k > 0:
        max_scores, _ = scores.max(-1)
        _, topk_inds = max_scores.topk(pre_top_k)
        batch_inds = torch.arange(batch_size).view(-1, 1).long()
        boxes = boxes[batch_inds, topk_inds, :]
        scores = scores[batch_inds, topk_inds, :]
    scores = scores.permute(0, 2, 1)
    selected_indices = ONNXNMSop.apply(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold)
    dets, labels = select_nms_index(scores, boxes, selected_indices, batch_size, keep_top_k=keep_top_k)
    return dets, labels


MARK_FUNCTION_COUNT = dict()


TORCH_DTYPE_TO_ONNX = {torch.uint8: torch.onnx.TensorProtoDataType.UINT8, torch.int8: torch.onnx.TensorProtoDataType.INT8, torch.float64: torch.onnx.TensorProtoDataType.DOUBLE, torch.float32: torch.onnx.TensorProtoDataType.FLOAT, torch.float16: torch.onnx.TensorProtoDataType.FLOAT16, torch.int32: torch.onnx.TensorProtoDataType.INT32, torch.int64: torch.onnx.TensorProtoDataType.INT64, torch.int16: torch.onnx.TensorProtoDataType.INT16, torch.bool: torch.onnx.TensorProtoDataType.BOOL, torch.complex64: torch.onnx.TensorProtoDataType.COMPLEX64, torch.complex128: torch.onnx.TensorProtoDataType.COMPLEX128}


class Mark(torch.autograd.Function):
    """The mark node function.

    The function does nothing but inserts a mark node to ONNX model. The mark
    can ease the process of model partition.
    """

    @staticmethod
    def symbolic(g, x, dtype, shape, func, func_id, type, name, id, attrs):
        """Symbolic function for mmdeploy::Mark op."""
        n = g.op('mmdeploy::Mark', x, dtype_i=TORCH_DTYPE_TO_ONNX[dtype], shape_i=shape, func_s=func, func_id_i=func_id, type_s=type, name_s=name, id_i=id, **attrs)
        return n

    @staticmethod
    def forward(ctx, x, *args) ->torch.Tensor:
        """Run forward."""
        return x


def mark_tensors(xs: Any, func: str, func_id: int, io_type: str, ctx: Any, attrs: Dict, is_inspecting: bool, level: int) ->tuple:
    """Add mark node recursively.

    Args:
        xs (Any): Input structure which contains tensor.
        func (str): Function name of the function which xs comes from.
        func_id (int): Function index of `func` in the model.
        io_type (str): The io type of xs, `input` or `output`.
        ctx (Any): The context instance.
        attrs (Dict): The extra attributes provided by mark decorator.
        is_inspecting (bool): The names of xs are inspected or not.
        level (int): The recursive level.

    Returns:
        Any: The same structure as xs, all tensor has been replaced with Mark.
    """
    visit = set()
    index = 0

    def impl(ys, prefix, level):
        nonlocal index
        old_index = index
        ret = ys
        prefix = () if level == 0 else prefix
        if isinstance(ys, torch.Tensor):
            if ys not in visit:
                visit.add(ys)
                root = ctx.names[ctx.index]
                name = '.'.join(str(x) for x in (root, *prefix))
                ys_shape = tuple(int(s) for s in ys.shape)
                ret = Mark.apply(ys, ys.dtype, ys_shape, func, func_id, io_type, name, index, attrs)
                index += 1
        elif isinstance(ys, list):
            ret = [impl(y, prefix + (i,), level + 1) for i, y in enumerate(ys)]
        elif isinstance(ys, tuple):
            ret = tuple(impl(y, prefix + (i,), level + 1) for i, y in enumerate(ys))
        elif isinstance(ys, dict):
            ret = {k: impl(v, prefix + (k,), level + 1) for k, v in ys.items()}
        if level == 0 and (is_inspecting or old_index != index):
            ctx.index += 1
        return ret
    return impl(xs, (), level)


def mark(func_name: Optional[str]=None, inputs: Optional[Sequence[str]]=None, outputs: Optional[Sequence[str]]=None, **attrs) ->Callable:
    """The decorator used to add mark node.

    Mark node can be used to support model partition.

    Args:
        func_name (str): The name of the function where marks come from.
        inputs (Sequence[str]): The input names of the marks. The final name
            might have suffix if inputs is list or dictionary.
        outputs (Sequence[str]): The output names of the marks. The final
            name might have suffix if outputs is list or dictionary.

    Returns:
        Callable: The process of mark decorator.

    Examples:
        >>> from mmdeploy.core import FUNCTION_REWRITER, mark
        >>> @FUNCTION_REWRITER.register_rewriter(
        >>>     func_name='mmdet.models.roi_heads.ConvFCBBoxHead.forward')
        >>> @mark(
        >>>     'bbox_head_forward',
        >>>     inputs=['bbox_feats'],
        >>>     outputs=['cls_score', 'bbox_pred'])
        >>> def forward_of_bbox_head(ctx, self, x):
        >>>     return ctx.origin_func(self, x)
    """
    MARK_FUNCTION_COUNT[func_name] = 0


    class Context:

        def __init__(self, names):
            self.names = names
            self.index = 0

    def decorator(f):
        func = func_name if func_name else f.__name__
        is_inspect = False
        if not inputs:
            input_names = list(inspect.signature(f).parameters.keys())
            is_inspect = True
        else:
            input_names = inputs
        output_names = outputs if outputs else func
        args_level, rets_level = -1, -1
        if isinstance(input_names, str):
            input_names = input_names,
        if isinstance(output_names, str):
            output_names = output_names,
            rets_level += 1

        def g(*args, **kwargs):
            func_id = MARK_FUNCTION_COUNT[func_name]
            MARK_FUNCTION_COUNT[func_name] += 1
            ctx = Context(input_names)
            args = mark_tensors(args, func, func_id, 'input', ctx, attrs, is_inspect, args_level)
            rets = f(*args, **kwargs)
            ctx = Context(output_names)
            func_ret = mark_tensors(rets, func, func_id, 'output', ctx, attrs, False, rets_level)
            return func_ret
        return g
    return decorator


@mark('multiclass_nms', inputs=['boxes', 'scores'], outputs=['dets', 'labels'])
def multiclass_nms(*args, **kwargs):
    """Wrapper function for `_multiclass_nms`."""
    return _multiclass_nms(*args, **kwargs)


class Checker(metaclass=ABCMeta):
    """The interface for checking whether a rewriter is valid."""

    def __init__(self):
        pass

    @abstractmethod
    def check(self, env: Dict) ->bool:
        """Check the if the rewriter is valid according to environment.

        Args:
            env (Dict): The backend, IR info and version info.
        """
        pass


class ContextCaller:
    """A callable object used in RewriteContext.

    This class saves context variables as member variables. When a rewritten
    function is called in RewriteContext, an instance of this class will be
    passed as the first argument of the function.

    Args:
        func (Callable): The rewritten function to call.
        origin_func (Callable): The function that is going to be rewritten.
            Note that in symbolic function origin_func may be 'None'.
        cfg (Dict): The deploy config dictionary.

    Example:
        >>> @FUNCTION_REWRITER.register_rewriter(func_name='torch.add')
        >>> def func(ctx, x, y):
        >>>     # ctx is an instance of ContextCaller
        >>>     print(ctx.cfg)
        >>>     return x + y
    """

    def __init__(self, func: Callable, origin_func: Callable, cfg: Dict, **kwargs):
        self.func = func
        self.origin_func = origin_func
        self.cfg = cfg
        if origin_func is not None:
            wraps(origin_func)(self)
        else:
            self.__annotations__ = getattr(func, '__annotations__', {})
        for k, v in kwargs.items():
            setattr(self, k, v)

    def __call__(self, *args, **kwargs):
        """Directly call self.func."""
        return self.func(self, *args, **kwargs)

    def get_wrapped_caller(self):
        """Generate a wrapped caller for function rewrite."""

        def wrapper(*args, **kwargs):
            return self.func(self, *args, **kwargs)
        return wrapper


def _del_func(path: str):
    """Delete a function that is denoted by a path.

    Args:
        path (str): The path to evaluate.
    """
    split_path = path.split('.')
    for i in range(len(split_path), 0, -1):
        try:
            exec('import {}'.format('.'.join(split_path[:i])))
            exec(f'del {path}')
            break
        except Exception:
            continue


def _replace_all_obj(obj: Any, new_obj: Any, ignore_refs: Tuple[Any]=tuple(), ignore_keys: Tuple[str]=tuple()):
    """Replace all object reference with new_object.

    Args:
        obj (Any): The object to be replaced.
        new_obj (Any): The object to replace obj.
        ignore_refs (Tuple[Any]): These refs will be ignored.
        ignore_keys (Tuple[str]): object with these keys will be ignored.
    """
    refs = gc.get_referrers(obj)
    obj_id = id(obj)
    for ref in refs:
        if ref in ignore_refs:
            continue
        elif isinstance(ref, MutableSequence):
            for i, v in enumerate(ref):
                if id(v) == obj_id:
                    ref[i] == new_obj
        elif isinstance(ref, Dict):
            for k, v in ref.items():
                if id(v) == obj_id and k not in ignore_keys:
                    ref[k] = new_obj
        else:
            pass


def _set_func(origin_func_path: str, rewrite_func: Callable, ignore_refs: Tuple[Any]=tuple(), ignore_keys: Tuple[str]=('origin_func',)):
    """Rewrite a function by executing a python statement.

    Args:
        origin_func_path (str): The path to origin function.
        rewrite_func (Callable): The new function instance.
        ignore_refs (Tuple[Any]): These refs will be ignored.
        ignore_keys (Tuple[str]): object with these keys will be ignored.
    """
    split_path = origin_func_path.split('.')
    for i in range(len(split_path), 0, -1):
        try:
            exec('import {}'.format('.'.join(split_path[:i])))
            break
        except Exception:
            continue
    origin_func = eval(origin_func_path)
    method_class = False
    if len(split_path) > 1:
        module_or_class = eval('.'.join(split_path[:-1]))
        if isinstance(module_or_class, type):
            method_class = True
    if not method_class:
        _replace_all_obj(origin_func, rewrite_func, ignore_refs=ignore_refs, ignore_keys=ignore_keys)
    exec(f'{origin_func_path} = rewrite_func')


def import_function(path: str) ->Tuple[Callable, Optional[type]]:
    """Import and evaluate a function. If the function is defined in a class,
    evaluate the class additionally.

    Args:
        path (str): The path to evaluate.

    Returns:
        Callable: The function of evaluation.
        type: The class of evaluation if the function is defined in a class, or
            None.
    """
    split_path = path.split('.')
    for i in range(len(split_path), 0, -1):
        try:
            exec('import {}'.format('.'.join(split_path[:i])))
            break
        except Exception:
            continue
    obj = eval(path)
    previous_obj = eval('.'.join(split_path[:-1]))
    if inspect.isclass(previous_obj):
        return obj, previous_obj
    else:
        return obj, None


def eval_with_import(path: str) ->Any:
    """Evaluate the string as Python script.

    Args:
        path (str): The path to evaluate.

    Returns:
        Any: The result of evaluation.
    """
    split_path = path.split('.')
    for i in range(len(split_path), 0, -1):
        try:
            exec('import {}'.format('.'.join(split_path[:i])))
            break
        except Exception:
            continue
    return eval(path)


class RewriterManager:
    """The rewrite manager that manages some rewriters."""

    def __init__(self):
        self.module_rewriter = ModuleRewriter()
        self.function_rewriter = FunctionRewriter()
        self.symbolic_rewriter = SymbolicRewriter()


class SARNet(nn.Module):
    """SARNet network structure for OCR image recognition.

    SARNet apply an argument named `valid_ratio` in network structure and this
    specific argument can not be ignored or replaced during exporting to onnx.
    Hence, some information of model config are utilized to compute
    `valid_ratio`.

    Paper: Show, attend and read: A simple and strong baseline for irregular
    text recognition.
    """

    def __init__(self, module, deploy_cfg, **kwargs):
        super(SARNet, self).__init__()
        self._module = module
        self.deploy_cfg = deploy_cfg

    def forward(self, img, *args, **kwargs):
        """Run forward."""
        img_metas = [{}]
        assert utils.is_type_list(img_metas, dict)
        assert isinstance(img, torch.Tensor)
        is_dynamic_flag = is_dynamic_shape(self.deploy_cfg)
        img_shape = torch._shape_as_tensor(img)[2:]
        if not is_dynamic_flag:
            img_shape = [int(val) for val in img_shape]
        img_metas[0]['img_shape'] = img_shape
        min_width, max_width, kepp_aspect_ratio = get_resize_ocr(self._module.cfg)
        valid_ratio = torch.tensor(1.0, device=img.device)
        if kepp_aspect_ratio:
            valid_ratio = 1 - (img == -1).sum() / np.array(img.size()).prod()
        valid_ratio = torch.clamp(valid_ratio, min_width / max_width, 1.0)
        img_metas[0]['valid_ratio'] = valid_ratio
        return self._module.simple_test(img, img_metas, *args, **kwargs)


class SequentialSARDecoder(nn.Module):
    """Rewrite Sequential Decoder module in `SAR.

    SequentialSARDecoder apply nn.LSTMCell inside, which brings obstacles to
    deployment. LSTMCell can be only exported to onnx in cpu device. To make it
    exportable to gpu device, LSTM is used to replace LSTMCell.

    <https://arxiv.org/abs/1811.00751>`_.
    """

    def __init__(self, module, deploy_cfg, **kwargs):
        super(SequentialSARDecoder, self).__init__()

        def lstmcell2lstm_params(lstm_mod, lstmcell_mod):
            lstm_mod.weight_ih_l0 = nn.Parameter(lstmcell_mod.weight_ih)
            lstm_mod.weight_hh_l0 = nn.Parameter(lstmcell_mod.weight_hh)
            lstm_mod.bias_ih_l0 = nn.Parameter(lstmcell_mod.bias_ih)
            lstm_mod.bias_hh_l0 = nn.Parameter(lstmcell_mod.bias_hh)
        self._module = module
        self.deploy_cfg = deploy_cfg
        if not self._module.dec_gru:
            rnn_decoder_layer1 = copy.deepcopy(self._module.rnn_decoder_layer1)
            rnn_decoder_layer2 = copy.deepcopy(self._module.rnn_decoder_layer2)
            self._module.rnn_decoder_layer1 = nn.LSTM(rnn_decoder_layer1.input_size, rnn_decoder_layer1.hidden_size, 1)
            self._module.rnn_decoder_layer2 = nn.LSTM(rnn_decoder_layer2.input_size, rnn_decoder_layer2.hidden_size, 1)
            lstmcell2lstm_params(self._module.rnn_decoder_layer1, rnn_decoder_layer1)
            lstmcell2lstm_params(self._module.rnn_decoder_layer2, rnn_decoder_layer2)
        self._module.train_mode = False

    def forward(self, feat, out_enc, targets_dict=None, img_metas=None, train_mode=True):
        return self._module.forward_test(feat, out_enc, img_metas)


class WrapFunction(nn.Module):
    """Wrap a pytorch function to nn.Module.

    It serves as partial function and can be exportable to ONNX.

    Args:
        wrapped_function (Callable): Input function to be wrapped.

    Examples:
        >>> from mmdeploy.utils.test import WrapFunction
        >>> import torch
        >>>
        >>> def clip(x, min, max):
        >>>     return torch.clamp(x, min, max)
        >>>
        >>> wrapped_model = WrapFunction(clip)
    """

    def __init__(self, wrapped_function: Callable, **kwargs):
        super(WrapFunction, self).__init__()
        self.wrapped_function = wrapped_function
        self.kwargs = kwargs

    def forward(self, *args, **kwargs) ->Any:
        """Call the wrapped function."""
        kwargs.update(self.kwargs)
        return self.wrapped_function(*args, **kwargs)


class WrapModel(nn.Module):
    """A wrapper class for rewrite unittests.

    It serves as partial function but can be rewritten with RewriterContext.

    Args:
        model (nn.Module): A pytorch module.
        func_name (str): Which function to use as forward function.

    Examples:
        >>> from mmdeploy.utils.test import WrapModel
        >>> from mmdet.models import AnchorHead
        >>>
        >>> model = AnchorHead(num_classes=4, in_channels=1)
        >>> wrapped_model = WrapModel(anchor_head,
        >>>                           'get_bboxes',
        >>>                           with_nms=True)
    """

    def __init__(self, model: nn.Module, func_name: str, **kwargs):
        super(WrapModel, self).__init__()
        assert hasattr(model, func_name), f'Got unexpected func name: {func_name}'
        self.model = model
        self.kwargs = kwargs
        self.func_name = func_name

    def forward(self, *args, **kwargs):
        """Run forward of the model."""
        kwargs.update(self.kwargs)
        func = getattr(self.model, self.func_name)
        return func(*args, **kwargs)


class DummyModel(torch.nn.Module):
    """A dummy model for unit tests.

    Args:
        outputs (Any): Predefined output variables.
    """

    def __init__(self, outputs=None, *args, **kwargs):
        torch.nn.Module.__init__(self)
        self.outputs = outputs

    def forward(self, *args, **kwargs):
        """Run forward."""
        return self.outputs

    def __call__(self, *args, **kwds):
        """Call the forward method."""
        return self.forward(*args, **kwds)


test_img = torch.rand(1, 3, 8, 8)


class DummyWrapper(torch.nn.Module):

    def __init__(self, outputs):
        self.outputs = outputs

    def __call__(self, *arg, **kwargs):
        return 0

    def output_to_list(self, *arg, **kwargs):
        return self.outputs


class ExampleBackbone(nn.Module):

    def __init__(self):
        super(ExampleBackbone, self).__init__()
        self.conv = nn.Conv2d(3, 3, 3)

    def init_weights(self, pretrained=None):
        pass

    def forward(self, x):
        return [self.conv(x)]


class OpModel(torch.nn.Module):

    def __init__(self, func, *args):
        super().__init__()
        self._func = func
        self._arg_tuple = args

    def forward(self, x):
        return self._func(x, *self._arg_tuple)


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (DummyModel,
     lambda: ([], {}),
     lambda: ([], {}),
     False),
    (DummyWrapper,
     lambda: ([], {'outputs': 4}),
     lambda: ([], {}),
     False),
    (ExampleBackbone,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     True),
    (OpModel,
     lambda: ([], {'func': _mock_layer()}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (WrapFunction,
     lambda: ([], {'wrapped_function': _mock_layer()}),
     lambda: ([], {'input': torch.rand([4, 4])}),
     False),
]

class Test_open_mmlab_mmdeploy(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

    def test_004(self):
        self._check(*TESTCASES[4])

