import sys
_module = sys.modules[__name__]
del sys
gather_models = _module
gen_benchmark_script = _module
doc_link_checker = _module
coco_instance = _module
nuim_instance = _module
default_runtime = _module
cascade_mask_rcnn_r50_fpn = _module
centerpoint_01voxel_second_secfpn_nus = _module
centerpoint_02pillar_second_secfpn_nus = _module
dgcnn = _module
fcaf3d = _module
fcos3d = _module
groupfree3d = _module
h3dnet = _module
hv_pointpillars_fpn_lyft = _module
hv_pointpillars_fpn_nus = _module
hv_pointpillars_fpn_range100_lyft = _module
hv_pointpillars_secfpn_kitti = _module
hv_pointpillars_secfpn_waymo = _module
hv_second_secfpn_kitti = _module
hv_second_secfpn_waymo = _module
imvotenet_image = _module
mask_rcnn_r50_fpn = _module
paconv_cuda_ssg = _module
paconv_ssg = _module
parta2 = _module
pgd = _module
point_rcnn = _module
pointnet2_msg = _module
pointnet2_ssg = _module
smoke = _module
votenet = _module
cosine = _module
cyclic_20e = _module
cyclic_40e = _module
mmdet_schedule_1x = _module
schedule_2x = _module
schedule_3x = _module
seg_cosine_100e = _module
seg_cosine_150e = _module
seg_cosine_200e = _module
seg_cosine_50e = _module
centerpoint_0075voxel_second_secfpn_4x8_cyclic_20e_nus = _module
centerpoint_0075voxel_second_secfpn_circlenms_4x8_cyclic_20e_nus = _module
centerpoint_0075voxel_second_secfpn_dcn_4x8_cyclic_20e_nus = _module
centerpoint_0075voxel_second_secfpn_dcn_4x8_cyclic_tta_20e_nus = _module
centerpoint_0075voxel_second_secfpn_dcn_circlenms_4x8_cyclic_20e_nus = _module
centerpoint_01voxel_second_secfpn_4x8_cyclic_20e_nus = _module
centerpoint_01voxel_second_secfpn_circlenms_4x8_cyclic_20e_nus = _module
centerpoint_01voxel_second_secfpn_dcn_4x8_cyclic_20e_nus = _module
centerpoint_01voxel_second_secfpn_dcn_circlenms_4x8_cyclic_20e_nus = _module
centerpoint_02pillar_second_secfpn_4x8_cyclic_20e_nus = _module
centerpoint_02pillar_second_secfpn_4x8_cyclic_20e_nus_novelo = _module
centerpoint_02pillar_second_secfpn_circlenms_4x8_cyclic_20e_nus = _module
centerpoint_02pillar_second_secfpn_dcn_4x8_cyclic_20e_nus = _module
centerpoint_02pillar_second_secfpn_dcn_circlenms_4x8_cyclic_20e_nus = _module
cascade_mask_rcnn_r101_fpn_1x_nuim = _module
cascade_mask_rcnn_r50_fpn_1x_nuim = _module
cascade_mask_rcnn_x101_32x4d_fpn_1x_nuim = _module
htc_r50_fpn_1x_nuim = _module
htc_without_semantic_r50_fpn_1x_nuim = _module
mask_rcnn_r101_fpn_1x_nuim = _module
mask_rcnn_r50_caffe_fpn_1x_nuim = _module
mask_rcnn_r50_fpn_1x_nuim = _module
mask_rcnn_x101_32x4d_fpn_1x_nuim = _module
collect_indoor3d_data = _module
indoor3d_util = _module
batch_load_scannet_data = _module
extract_posed_images = _module
load_scannet_data = _module
scannet_utils = _module
mono_det_demo = _module
multi_modality_demo = _module
pc_seg_demo = _module
pcd_demo = _module
conf = _module
stat = _module
mmdet3d = _module
apis = _module
inference = _module
test = _module
train = _module
core = _module
anchor = _module
anchor_3d_generator = _module
bbox = _module
assigners = _module
box_np_ops = _module
coders = _module
anchor_free_bbox_coder = _module
centerpoint_bbox_coders = _module
delta_xyzwhlr_bbox_coder = _module
fcos3d_bbox_coder = _module
groupfree3d_bbox_coder = _module
monoflex_bbox_coder = _module
partial_bin_based_bbox_coder = _module
pgd_bbox_coder = _module
point_xyzwhlr_bbox_coder = _module
smoke_bbox_coder = _module
iou_calculators = _module
iou3d_calculator = _module
samplers = _module
iou_neg_piecewise_sampler = _module
structures = _module
base_box3d = _module
box_3d_mode = _module
cam_box3d = _module
coord_3d_mode = _module
depth_box3d = _module
lidar_box3d = _module
utils = _module
transforms = _module
evaluation = _module
indoor_eval = _module
instance_seg_eval = _module
kitti_utils = _module
eval = _module
rotate_iou = _module
lyft_eval = _module
evaluate_semantic_instance = _module
util_3d = _module
seg_eval = _module
waymo_utils = _module
prediction_kitti_to_waymo = _module
points = _module
base_points = _module
cam_points = _module
depth_points = _module
lidar_points = _module
post_processing = _module
box3d_nms = _module
merge_augs = _module
array_converter = _module
gaussian = _module
visualizer = _module
image_vis = _module
open3d_vis = _module
show_result = _module
voxel = _module
builder = _module
voxel_generator = _module
datasets = _module
custom_3d = _module
custom_3d_seg = _module
dataset_wrappers = _module
kitti2d_dataset = _module
kitti_dataset = _module
kitti_mono_dataset = _module
lyft_dataset = _module
nuscenes_dataset = _module
nuscenes_mono_dataset = _module
pipelines = _module
compose = _module
data_augment_utils = _module
dbsampler = _module
formating = _module
loading = _module
test_time_aug = _module
transforms_3d = _module
s3dis_dataset = _module
scannet_dataset = _module
semantickitti_dataset = _module
sunrgbd_dataset = _module
utils = _module
waymo_dataset = _module
models = _module
backbones = _module
base_pointnet = _module
dgcnn = _module
dla = _module
mink_resnet = _module
multi_backbone = _module
nostem_regnet = _module
pointnet2_sa_msg = _module
pointnet2_sa_ssg = _module
second = _module
decode_heads = _module
decode_head = _module
dgcnn_head = _module
paconv_head = _module
pointnet2_head = _module
dense_heads = _module
anchor3d_head = _module
anchor_free_mono3d_head = _module
base_conv_bbox_head = _module
base_mono3d_dense_head = _module
centerpoint_head = _module
fcaf3d_head = _module
fcos_mono3d_head = _module
free_anchor3d_head = _module
groupfree3d_head = _module
imvoxel_head = _module
monoflex_head = _module
parta2_rpn_head = _module
pgd_head = _module
point_rpn_head = _module
shape_aware_head = _module
smoke_mono3d_head = _module
ssd_3d_head = _module
train_mixins = _module
vote_head = _module
detectors = _module
base = _module
centerpoint = _module
dynamic_voxelnet = _module
fcos_mono3d = _module
groupfree3dnet = _module
h3dnet = _module
imvotenet = _module
imvoxelnet = _module
mink_single_stage = _module
mvx_faster_rcnn = _module
mvx_two_stage = _module
parta2 = _module
point_rcnn = _module
sassd = _module
single_stage = _module
single_stage_mono3d = _module
smoke_mono3d = _module
ssd3dnet = _module
two_stage = _module
votenet = _module
voxelnet = _module
fusion_layers = _module
coord_transform = _module
point_fusion = _module
vote_fusion = _module
losses = _module
axis_aligned_iou_loss = _module
chamfer_distance = _module
multibin_loss = _module
paconv_regularization_loss = _module
rotated_iou_loss = _module
uncertain_smooth_l1_loss = _module
middle_encoders = _module
pillar_scatter = _module
sparse_encoder = _module
sparse_unet = _module
model_utils = _module
edge_fusion_module = _module
transformer = _module
vote_module = _module
necks = _module
dla_neck = _module
imvoxel_neck = _module
pointnet2_fp_neck = _module
second_fpn = _module
view_transformer = _module
roi_heads = _module
base_3droi_head = _module
bbox_heads = _module
h3d_bbox_head = _module
parta2_bbox_head = _module
point_rcnn_bbox_head = _module
h3d_roi_head = _module
mask_heads = _module
pointwise_semantic_head = _module
primitive_head = _module
part_aggregation_roi_head = _module
point_rcnn_roi_head = _module
roi_extractors = _module
single_roiaware_extractor = _module
single_roipoint_extractor = _module
segmentors = _module
base = _module
encoder_decoder = _module
clip_sigmoid = _module
edge_indices = _module
gen_keypoints = _module
handle_objs = _module
mlp = _module
voxel_encoders = _module
pillar_encoder = _module
utils = _module
voxel_encoder = _module
ops = _module
dgcnn_modules = _module
dgcnn_fa_module = _module
dgcnn_fp_module = _module
dgcnn_gf_module = _module
norm = _module
paconv = _module
paconv = _module
utils = _module
pointnet_modules = _module
builder = _module
paconv_sa_module = _module
point_fp_module = _module
point_sa_module = _module
sparse_block = _module
spconv = _module
overwrite_spconv = _module
write_spconv2 = _module
collect_env = _module
compat_cfg = _module
logger = _module
misc = _module
setup_env = _module
version = _module
dummy = _module
dummy_resnet = _module
setup = _module
test_dataset_wrappers = _module
test_kitti_dataset = _module
test_kitti_mono_dataset = _module
test_lyft_dataset = _module
test_nuscene_dataset = _module
test_nuscenes_mono_dataset = _module
test_s3dis_dataset = _module
test_scannet_dataset = _module
test_semantickitti_dataset = _module
test_sunrgbd_dataset = _module
test_waymo_dataset = _module
test_data_augment_utils = _module
test_test_augment_utils = _module
test_transforms_3d = _module
test_indoor_pipeline = _module
test_indoor_sample = _module
test_load_images_from_multi_views = _module
test_load_points_from_multi_sweeps = _module
test_loading = _module
test_outdoor_pipeline = _module
test_indoor_eval = _module
test_instance_seg_eval = _module
test_kitti_eval = _module
test_losses = _module
test_seg_eval = _module
test_backbones = _module
test_dgcnn_modules = _module
test_middle_encoders = _module
test_paconv_modules = _module
test_paconv_ops = _module
test_pointnet_modules = _module
test_sparse_unet = _module
test_vote_module = _module
test_detectors = _module
test_forward = _module
test_fusion_coord_trans = _module
test_point_fusion = _module
test_vote_fusion = _module
test_dgcnn_decode_head = _module
test_heads = _module
test_paconv_decode_head = _module
test_parta2_bbox_head = _module
test_pointnet2_decode_head = _module
test_roi_extractors = _module
test_semantic_heads = _module
test_fpn = _module
test_necks = _module
test_segmentors = _module
test_voxel_encoders = _module
test_voxel_generator = _module
test_apis = _module
test_config = _module
test_anchors = _module
test_assigners = _module
test_bbox_coders = _module
test_box3d = _module
test_box_np_ops = _module
test_compat_cfg = _module
test_coord_3d_mode = _module
test_merge_augs = _module
test_nms = _module
test_points = _module
test_samplers = _module
test_setup_env = _module
test_utils = _module
analyze_logs = _module
benchmark = _module
get_flops = _module
create_data = _module
data_converter = _module
create_gt_database = _module
indoor_converter = _module
kitti_converter = _module
kitti_data_utils = _module
lyft_converter = _module
lyft_data_fixer = _module
nuimage_converter = _module
nuscenes_converter = _module
s3dis_data_utils = _module
scannet_data_utils = _module
sunrgbd_data_utils = _module
waymo_converter = _module
mmdet3d2torchserve = _module
mmdet3d_handler = _module
test_torchserver = _module
browse_dataset = _module
fuse_conv_bn = _module
print_config = _module
visualize_results = _module
convert_h3dnet_checkpoints = _module
convert_votenet_checkpoints = _module
publish_model = _module
regnet2mmdet = _module
test = _module
train = _module
update_data_coords = _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


import re


from copy import deepcopy


import numpy as np


import random


import warnings


from torch import distributed as dist


from torch.nn import functional as F


import math


from abc import abstractmethod


from enum import IntEnum


from enum import unique


from logging import warning


import functools


from inspect import getfullargspec


import copy


from matplotlib import pyplot as plt


from torch.utils.data import Dataset


from abc import ABCMeta


from torch import nn as nn


from torch import nn


import torch.nn as nn


from functools import partial


from torch.nn.functional import l1_loss


from torch.nn.functional import mse_loss


from torch.nn.functional import smooth_l1_loss


from torch.autograd.function import Function


from typing import List


import itertools


from torch.nn.parameter import Parameter


from torch import multiprocessing as mp


from torch.utils.cpp_extension import BuildExtension


from torch.utils.cpp_extension import CppExtension


from torch.utils.cpp_extension import CUDAExtension


from torch.nn import BatchNorm1d


from torch.nn import ReLU


import time


from collections import OrderedDict


import torch.distributed as dist


def dla_build_norm_layer(cfg, num_features):
    """Build normalization layer specially designed for DLANet.

    Args:
        cfg (dict): The norm layer config, which should contain:

            - type (str): Layer type.
            - layer args: Args needed to instantiate a norm layer.
            - requires_grad (bool, optional): Whether stop gradient updates.
        num_features (int): Number of input channels.


    Returns:
        Function: Build normalization layer in mmcv.
    """
    cfg_ = cfg.copy()
    if cfg_['type'] == 'GN':
        if num_features % 32 == 0:
            return build_norm_layer(cfg_, num_features)
        else:
            assert 'num_groups' in cfg_
            cfg_['num_groups'] = cfg_['num_groups'] // 2
            return build_norm_layer(cfg_, num_features)
    else:
        return build_norm_layer(cfg_, num_features)


class GeneralSamplingModule(nn.Module):
    """Sampling Points.

    Sampling points with given index.
    """

    def forward(self, xyz, features, sample_inds):
        """Forward pass.

        Args:
            xyz： (B, N, 3) the coordinates of the features.
            features (Tensor): (B, C, N) features to sample.
            sample_inds (Tensor): (B, M) the given index,
                where M is the number of points.

        Returns:
            Tensor: (B, M, 3) coordinates of sampled features
            Tensor: (B, C, M) the sampled features.
            Tensor: (B, M) the given index.
        """
        xyz_t = xyz.transpose(1, 2).contiguous()
        new_xyz = gather_points(xyz_t, sample_inds).transpose(1, 2).contiguous()
        new_features = gather_points(features, sample_inds).contiguous()
        return new_xyz, new_features, sample_inds


EPS = 1e-06


class ArrayConverter:
    SUPPORTED_NON_ARRAY_TYPES = int, float, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64, np.float16, np.float32, np.float64

    def __init__(self, template_array=None):
        if template_array is not None:
            self.set_template(template_array)

    def set_template(self, array):
        """Set template array.

        Args:
            array (tuple | list | int | float | np.ndarray | torch.Tensor):
                Template array.

        Raises:
            ValueError: If input is list or tuple and cannot be converted to
                to a NumPy array, a ValueError is raised.
            TypeError: If input type does not belong to the above range,
                or the contents of a list or tuple do not share the
                same data type, a TypeError is raised.
        """
        self.array_type = type(array)
        self.is_num = False
        self.device = 'cpu'
        if isinstance(array, np.ndarray):
            self.dtype = array.dtype
        elif isinstance(array, torch.Tensor):
            self.dtype = array.dtype
            self.device = array.device
        elif isinstance(array, (list, tuple)):
            try:
                array = np.array(array)
                if array.dtype not in self.SUPPORTED_NON_ARRAY_TYPES:
                    raise TypeError
                self.dtype = array.dtype
            except (ValueError, TypeError):
                None
                raise
        elif isinstance(array, self.SUPPORTED_NON_ARRAY_TYPES):
            self.array_type = np.ndarray
            self.is_num = True
            self.dtype = np.dtype(type(array))
        else:
            raise TypeError(f'Template type {self.array_type} is not supported.')

    def convert(self, input_array, target_type=None, target_array=None):
        """Convert input array to target data type.

        Args:
            input_array (tuple | list | np.ndarray |
                torch.Tensor | int | float ):
                Input array. Defaults to None.
            target_type (<class 'np.ndarray'> | <class 'torch.Tensor'>,
                optional):
                Type to which input array is converted. Defaults to None.
            target_array (np.ndarray | torch.Tensor, optional):
                Template array to which input array is converted.
                Defaults to None.

        Raises:
            ValueError: If input is list or tuple and cannot be converted to
                to a NumPy array, a ValueError is raised.
            TypeError: If input type does not belong to the above range,
                or the contents of a list or tuple do not share the
                same data type, a TypeError is raised.
        """
        if isinstance(input_array, (list, tuple)):
            try:
                input_array = np.array(input_array)
                if input_array.dtype not in self.SUPPORTED_NON_ARRAY_TYPES:
                    raise TypeError
            except (ValueError, TypeError):
                None
                raise
        elif isinstance(input_array, self.SUPPORTED_NON_ARRAY_TYPES):
            input_array = np.array(input_array)
        array_type = type(input_array)
        assert target_type is not None or target_array is not None, 'must specify a target'
        if target_type is not None:
            assert target_type in (np.ndarray, torch.Tensor), 'invalid target type'
            if target_type == array_type:
                return input_array
            elif target_type == np.ndarray:
                converted_array = input_array.cpu().numpy().astype(np.float32)
            else:
                converted_array = torch.tensor(input_array, dtype=torch.float32)
        else:
            assert isinstance(target_array, (np.ndarray, torch.Tensor)), 'invalid target array type'
            if isinstance(target_array, array_type):
                return input_array
            elif isinstance(target_array, np.ndarray):
                converted_array = input_array.cpu().numpy().astype(target_array.dtype)
            else:
                converted_array = target_array.new_tensor(input_array)
        return converted_array

    def recover(self, input_array):
        assert isinstance(input_array, (np.ndarray, torch.Tensor)), 'invalid input array type'
        if isinstance(input_array, self.array_type):
            return input_array
        elif isinstance(input_array, torch.Tensor):
            converted_array = input_array.cpu().numpy().astype(self.dtype)
        else:
            converted_array = torch.tensor(input_array, dtype=self.dtype, device=self.device)
        if self.is_num:
            converted_array = converted_array.item()
        return converted_array


def array_converter(to_torch=True, apply_to=tuple(), template_arg_name_=None, recover=True):
    """Wrapper function for data-type agnostic processing.

    First converts input arrays to PyTorch tensors or NumPy ndarrays
    for middle calculation, then convert output to original data-type if
    `recover=True`.

    Args:
        to_torch (Bool, optional): Whether convert to PyTorch tensors
            for middle calculation. Defaults to True.
        apply_to (tuple[str], optional): The arguments to which we apply
            data-type conversion. Defaults to an empty tuple.
        template_arg_name_ (str, optional): Argument serving as the template (
            return arrays should have the same dtype and device
            as the template). Defaults to None. If None, we will use the
            first argument in `apply_to` as the template argument.
        recover (Bool, optional): Whether or not recover the wrapped function
            outputs to the `template_arg_name_` type. Defaults to True.

    Raises:
        ValueError: When template_arg_name_ is not among all args, or
            when apply_to contains an arg which is not among all args,
            a ValueError will be raised. When the template argument or
            an argument to convert is a list or tuple, and cannot be
            converted to a NumPy array, a ValueError will be raised.
        TypeError: When the type of the template argument or
                an argument to convert does not belong to the above range,
                or the contents of such an list-or-tuple-type argument
                do not share the same data type, a TypeError is raised.

    Returns:
        (function): wrapped function.

    Example:
        >>> import torch
        >>> import numpy as np
        >>>
        >>> # Use torch addition for a + b,
        >>> # and convert return values to the type of a
        >>> @array_converter(apply_to=('a', 'b'))
        >>> def simple_add(a, b):
        >>>     return a + b
        >>>
        >>> a = np.array([1.1])
        >>> b = np.array([2.2])
        >>> simple_add(a, b)
        >>>
        >>> # Use numpy addition for a + b,
        >>> # and convert return values to the type of b
        >>> @array_converter(to_torch=False, apply_to=('a', 'b'),
        >>>                  template_arg_name_='b')
        >>> def simple_add(a, b):
        >>>     return a + b
        >>>
        >>> simple_add()
        >>>
        >>> # Use torch funcs for floor(a) if flag=True else ceil(a),
        >>> # and return the torch tensor
        >>> @array_converter(apply_to=('a',), recover=False)
        >>> def floor_or_ceil(a, flag=True):
        >>>     return torch.floor(a) if flag else torch.ceil(a)
        >>>
        >>> floor_or_ceil(a, flag=False)
    """

    def array_converter_wrapper(func):
        """Outer wrapper for the function."""

        @functools.wraps(func)
        def new_func(*args, **kwargs):
            """Inner wrapper for the arguments."""
            if len(apply_to) == 0:
                return func(*args, **kwargs)
            func_name = func.__name__
            arg_spec = getfullargspec(func)
            arg_names = arg_spec.args
            arg_num = len(arg_names)
            default_arg_values = arg_spec.defaults
            if default_arg_values is None:
                default_arg_values = []
            no_default_arg_num = len(arg_names) - len(default_arg_values)
            kwonly_arg_names = arg_spec.kwonlyargs
            kwonly_default_arg_values = arg_spec.kwonlydefaults
            if kwonly_default_arg_values is None:
                kwonly_default_arg_values = {}
            all_arg_names = arg_names + kwonly_arg_names
            if len(args) > arg_num:
                named_args = args[:arg_num]
                nameless_args = args[arg_num:]
            else:
                named_args = args
                nameless_args = []
            if template_arg_name_ is None:
                template_arg_name = apply_to[0]
            else:
                template_arg_name = template_arg_name_
            if template_arg_name not in all_arg_names:
                raise ValueError(f'{template_arg_name} is not among the argument list of function {func_name}')
            for arg_to_apply in apply_to:
                if arg_to_apply not in all_arg_names:
                    raise ValueError(f'{arg_to_apply} is not an argument of {func_name}')
            new_args = []
            new_kwargs = {}
            converter = ArrayConverter()
            target_type = torch.Tensor if to_torch else np.ndarray
            for i, arg_value in enumerate(named_args):
                if arg_names[i] in apply_to:
                    new_args.append(converter.convert(input_array=arg_value, target_type=target_type))
                else:
                    new_args.append(arg_value)
                if arg_names[i] == template_arg_name:
                    template_arg_value = arg_value
            kwonly_default_arg_values.update(kwargs)
            kwargs = kwonly_default_arg_values
            for i in range(len(named_args), len(all_arg_names)):
                arg_name = all_arg_names[i]
                if arg_name in kwargs:
                    if arg_name in apply_to:
                        new_kwargs[arg_name] = converter.convert(input_array=kwargs[arg_name], target_type=target_type)
                    else:
                        new_kwargs[arg_name] = kwargs[arg_name]
                else:
                    default_value = default_arg_values[i - no_default_arg_num]
                    if arg_name in apply_to:
                        new_kwargs[arg_name] = converter.convert(input_array=default_value, target_type=target_type)
                    else:
                        new_kwargs[arg_name] = default_value
                if arg_name == template_arg_name:
                    template_arg_value = kwargs[arg_name]
            new_args += nameless_args
            return_values = func(*new_args, **new_kwargs)
            converter.set_template(template_arg_value)

            def recursive_recover(input_data):
                if isinstance(input_data, (tuple, list)):
                    new_data = []
                    for item in input_data:
                        new_data.append(recursive_recover(item))
                    return tuple(new_data) if isinstance(input_data, tuple) else new_data
                elif isinstance(input_data, dict):
                    new_data = {}
                    for k, v in input_data.items():
                        new_data[k] = recursive_recover(v)
                    return new_data
                elif isinstance(input_data, (torch.Tensor, np.ndarray)):
                    return converter.recover(input_data)
                else:
                    return input_data
            if recover:
                return recursive_recover(return_values)
            else:
                return return_values
        return new_func
    return array_converter_wrapper


@array_converter(apply_to=('points', 'angles'))
def rotation_3d_in_axis(points, angles, axis=0, return_mat=False, clockwise=False):
    """Rotate points by angles according to axis.

    Args:
        points (np.ndarray | torch.Tensor | list | tuple ):
            Points of shape (N, M, 3).
        angles (np.ndarray | torch.Tensor | list | tuple | float):
            Vector of angles in shape (N,)
        axis (int, optional): The axis to be rotated. Defaults to 0.
        return_mat: Whether or not return the rotation matrix (transposed).
            Defaults to False.
        clockwise: Whether the rotation is clockwise. Defaults to False.

    Raises:
        ValueError: when the axis is not in range [0, 1, 2], it will
            raise value error.

    Returns:
        (torch.Tensor | np.ndarray): Rotated points in shape (N, M, 3).
    """
    batch_free = len(points.shape) == 2
    if batch_free:
        points = points[None]
    if isinstance(angles, float) or len(angles.shape) == 0:
        angles = torch.full(points.shape[:1], angles)
    assert len(points.shape) == 3 and len(angles.shape) == 1 and points.shape[0] == angles.shape[0], f'Incorrect shape of points angles: {points.shape}, {angles.shape}'
    assert points.shape[-1] in [2, 3], f'Points size should be 2 or 3 instead of {points.shape[-1]}'
    rot_sin = torch.sin(angles)
    rot_cos = torch.cos(angles)
    ones = torch.ones_like(rot_cos)
    zeros = torch.zeros_like(rot_cos)
    if points.shape[-1] == 3:
        if axis == 1 or axis == -2:
            rot_mat_T = torch.stack([torch.stack([rot_cos, zeros, -rot_sin]), torch.stack([zeros, ones, zeros]), torch.stack([rot_sin, zeros, rot_cos])])
        elif axis == 2 or axis == -1:
            rot_mat_T = torch.stack([torch.stack([rot_cos, rot_sin, zeros]), torch.stack([-rot_sin, rot_cos, zeros]), torch.stack([zeros, zeros, ones])])
        elif axis == 0 or axis == -3:
            rot_mat_T = torch.stack([torch.stack([ones, zeros, zeros]), torch.stack([zeros, rot_cos, rot_sin]), torch.stack([zeros, -rot_sin, rot_cos])])
        else:
            raise ValueError(f'axis should in range [-3, -2, -1, 0, 1, 2], got {axis}')
    else:
        rot_mat_T = torch.stack([torch.stack([rot_cos, rot_sin]), torch.stack([-rot_sin, rot_cos])])
    if clockwise:
        rot_mat_T = rot_mat_T.transpose(0, 1)
    if points.shape[0] == 0:
        points_new = points
    else:
        points_new = torch.einsum('aij,jka->aik', points, rot_mat_T)
    if batch_free:
        points_new = points_new.squeeze(0)
    if return_mat:
        rot_mat_T = torch.einsum('jka->ajk', rot_mat_T)
        if batch_free:
            rot_mat_T = rot_mat_T.squeeze(0)
        return points_new, rot_mat_T
    else:
        return points_new


class BasePoints(object):
    """Base class for Points.

    Args:
        tensor (torch.Tensor | np.ndarray | list): a N x points_dim matrix.
        points_dim (int, optional): Number of the dimension of a point.
            Each row is (x, y, z). Defaults to 3.
        attribute_dims (dict, optional): Dictionary to indicate the
            meaning of extra dimension. Defaults to None.

    Attributes:
        tensor (torch.Tensor): Float matrix of N x points_dim.
        points_dim (int): Integer indicating the dimension of a point.
            Each row is (x, y, z, ...).
        attribute_dims (bool): Dictionary to indicate the meaning of extra
            dimension. Defaults to None.
        rotation_axis (int): Default rotation axis for points rotation.
    """

    def __init__(self, tensor, points_dim=3, attribute_dims=None):
        if isinstance(tensor, torch.Tensor):
            device = tensor.device
        else:
            device = torch.device('cpu')
        tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)
        if tensor.numel() == 0:
            tensor = tensor.reshape((0, points_dim))
        assert tensor.dim() == 2 and tensor.size(-1) == points_dim, tensor.size()
        self.tensor = tensor
        self.points_dim = points_dim
        self.attribute_dims = attribute_dims
        self.rotation_axis = 0

    @property
    def coord(self):
        """torch.Tensor: Coordinates of each point in shape (N, 3)."""
        return self.tensor[:, :3]

    @coord.setter
    def coord(self, tensor):
        """Set the coordinates of each point."""
        try:
            tensor = tensor.reshape(self.shape[0], 3)
        except (RuntimeError, ValueError):
            raise ValueError(f'got unexpected shape {tensor.shape}')
        if not isinstance(tensor, torch.Tensor):
            tensor = self.tensor.new_tensor(tensor)
        self.tensor[:, :3] = tensor

    @property
    def height(self):
        """torch.Tensor:
            A vector with height of each point in shape (N, 1), or None."""
        if self.attribute_dims is not None and 'height' in self.attribute_dims.keys():
            return self.tensor[:, self.attribute_dims['height']]
        else:
            return None

    @height.setter
    def height(self, tensor):
        """Set the height of each point."""
        try:
            tensor = tensor.reshape(self.shape[0])
        except (RuntimeError, ValueError):
            raise ValueError(f'got unexpected shape {tensor.shape}')
        if not isinstance(tensor, torch.Tensor):
            tensor = self.tensor.new_tensor(tensor)
        if self.attribute_dims is not None and 'height' in self.attribute_dims.keys():
            self.tensor[:, self.attribute_dims['height']] = tensor
        else:
            if self.attribute_dims is None:
                self.attribute_dims = dict()
            attr_dim = self.shape[1]
            self.tensor = torch.cat([self.tensor, tensor.unsqueeze(1)], dim=1)
            self.attribute_dims.update(dict(height=attr_dim))
            self.points_dim += 1

    @property
    def color(self):
        """torch.Tensor:
            A vector with color of each point in shape (N, 3), or None."""
        if self.attribute_dims is not None and 'color' in self.attribute_dims.keys():
            return self.tensor[:, self.attribute_dims['color']]
        else:
            return None

    @color.setter
    def color(self, tensor):
        """Set the color of each point."""
        try:
            tensor = tensor.reshape(self.shape[0], 3)
        except (RuntimeError, ValueError):
            raise ValueError(f'got unexpected shape {tensor.shape}')
        if tensor.max() >= 256 or tensor.min() < 0:
            warnings.warn('point got color value beyond [0, 255]')
        if not isinstance(tensor, torch.Tensor):
            tensor = self.tensor.new_tensor(tensor)
        if self.attribute_dims is not None and 'color' in self.attribute_dims.keys():
            self.tensor[:, self.attribute_dims['color']] = tensor
        else:
            if self.attribute_dims is None:
                self.attribute_dims = dict()
            attr_dim = self.shape[1]
            self.tensor = torch.cat([self.tensor, tensor], dim=1)
            self.attribute_dims.update(dict(color=[attr_dim, attr_dim + 1, attr_dim + 2]))
            self.points_dim += 3

    @property
    def shape(self):
        """torch.Shape: Shape of points."""
        return self.tensor.shape

    def shuffle(self):
        """Shuffle the points.

        Returns:
            torch.Tensor: The shuffled index.
        """
        idx = torch.randperm(self.__len__(), device=self.tensor.device)
        self.tensor = self.tensor[idx]
        return idx

    def rotate(self, rotation, axis=None):
        """Rotate points with the given rotation matrix or angle.

        Args:
            rotation (float | np.ndarray | torch.Tensor): Rotation matrix
                or angle.
            axis (int, optional): Axis to rotate at. Defaults to None.
        """
        if not isinstance(rotation, torch.Tensor):
            rotation = self.tensor.new_tensor(rotation)
        assert rotation.shape == torch.Size([3, 3]) or rotation.numel() == 1, f'invalid rotation shape {rotation.shape}'
        if axis is None:
            axis = self.rotation_axis
        if rotation.numel() == 1:
            rotated_points, rot_mat_T = rotation_3d_in_axis(self.tensor[:, :3][None], rotation, axis=axis, return_mat=True)
            self.tensor[:, :3] = rotated_points.squeeze(0)
            rot_mat_T = rot_mat_T.squeeze(0)
        else:
            self.tensor[:, :3] = self.tensor[:, :3] @ rotation
            rot_mat_T = rotation
        return rot_mat_T

    @abstractmethod
    def flip(self, bev_direction='horizontal'):
        """Flip the points along given BEV direction.

        Args:
            bev_direction (str): Flip direction (horizontal or vertical).
        """
        pass

    def translate(self, trans_vector):
        """Translate points with the given translation vector.

        Args:
            trans_vector (np.ndarray, torch.Tensor): Translation
                vector of size 3 or nx3.
        """
        if not isinstance(trans_vector, torch.Tensor):
            trans_vector = self.tensor.new_tensor(trans_vector)
        trans_vector = trans_vector.squeeze(0)
        if trans_vector.dim() == 1:
            assert trans_vector.shape[0] == 3
        elif trans_vector.dim() == 2:
            assert trans_vector.shape[0] == self.tensor.shape[0] and trans_vector.shape[1] == 3
        else:
            raise NotImplementedError(f'Unsupported translation vector of shape {trans_vector.shape}')
        self.tensor[:, :3] += trans_vector

    def in_range_3d(self, point_range):
        """Check whether the points are in the given range.

        Args:
            point_range (list | torch.Tensor): The range of point
                (x_min, y_min, z_min, x_max, y_max, z_max)

        Note:
            In the original implementation of SECOND, checking whether
            a box in the range checks whether the points are in a convex
            polygon, we try to reduce the burden for simpler cases.

        Returns:
            torch.Tensor: A binary vector indicating whether each point is
                inside the reference range.
        """
        in_range_flags = (self.tensor[:, 0] > point_range[0]) & (self.tensor[:, 1] > point_range[1]) & (self.tensor[:, 2] > point_range[2]) & (self.tensor[:, 0] < point_range[3]) & (self.tensor[:, 1] < point_range[4]) & (self.tensor[:, 2] < point_range[5])
        return in_range_flags

    @property
    def bev(self):
        """torch.Tensor: BEV of the points in shape (N, 2)."""
        return self.tensor[:, [0, 1]]

    def in_range_bev(self, point_range):
        """Check whether the points are in the given range.

        Args:
            point_range (list | torch.Tensor): The range of point
                in order of (x_min, y_min, x_max, y_max).

        Returns:
            torch.Tensor: Indicating whether each point is inside
                the reference range.
        """
        in_range_flags = (self.bev[:, 0] > point_range[0]) & (self.bev[:, 1] > point_range[1]) & (self.bev[:, 0] < point_range[2]) & (self.bev[:, 1] < point_range[3])
        return in_range_flags

    @abstractmethod
    def convert_to(self, dst, rt_mat=None):
        """Convert self to ``dst`` mode.

        Args:
            dst (:obj:`CoordMode`): The target Box mode.
            rt_mat (np.ndarray | torch.Tensor, optional): The rotation and
                translation matrix between different coordinates.
                Defaults to None.
                The conversion from `src` coordinates to `dst` coordinates
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
            :obj:`BasePoints`: The converted box of the same type
                in the `dst` mode.
        """
        pass

    def scale(self, scale_factor):
        """Scale the points with horizontal and vertical scaling factors.

        Args:
            scale_factors (float): Scale factors to scale the points.
        """
        self.tensor[:, :3] *= scale_factor

    def __getitem__(self, item):
        """
        Note:
            The following usage are allowed:
            1. `new_points = points[3]`:
                return a `Points` that contains only one point.
            2. `new_points = points[2:10]`:
                return a slice of points.
            3. `new_points = points[vector]`:
                where vector is a torch.BoolTensor with `length = len(points)`.
                Nonzero elements in the vector will be selected.
            4. `new_points = points[3:11, vector]`:
                return a slice of points and attribute dims.
            5. `new_points = points[4:12, 2]`:
                return a slice of points with single attribute.
            Note that the returned Points might share storage with this Points,
            subject to Pytorch's indexing semantics.

        Returns:
            :obj:`BasePoints`: A new object of
                :class:`BasePoints` after indexing.
        """
        original_type = type(self)
        if isinstance(item, int):
            return original_type(self.tensor[item].view(1, -1), points_dim=self.points_dim, attribute_dims=self.attribute_dims)
        elif isinstance(item, tuple) and len(item) == 2:
            if isinstance(item[1], slice):
                start = 0 if item[1].start is None else item[1].start
                stop = self.tensor.shape[1] if item[1].stop is None else item[1].stop
                step = 1 if item[1].step is None else item[1].step
                item = list(item)
                item[1] = list(range(start, stop, step))
                item = tuple(item)
            elif isinstance(item[1], int):
                item = list(item)
                item[1] = [item[1]]
                item = tuple(item)
            p = self.tensor[item[0], item[1]]
            keep_dims = list(set(item[1]).intersection(set(range(3, self.tensor.shape[1]))))
            if self.attribute_dims is not None:
                attribute_dims = self.attribute_dims.copy()
                for key in self.attribute_dims.keys():
                    cur_attribute_dims = attribute_dims[key]
                    if isinstance(cur_attribute_dims, int):
                        cur_attribute_dims = [cur_attribute_dims]
                    intersect_attr = list(set(cur_attribute_dims).intersection(set(keep_dims)))
                    if len(intersect_attr) == 1:
                        attribute_dims[key] = intersect_attr[0]
                    elif len(intersect_attr) > 1:
                        attribute_dims[key] = intersect_attr
                    else:
                        attribute_dims.pop(key)
            else:
                attribute_dims = None
        elif isinstance(item, (slice, np.ndarray, torch.Tensor)):
            p = self.tensor[item]
            attribute_dims = self.attribute_dims
        else:
            raise NotImplementedError(f'Invalid slice {item}!')
        assert p.dim() == 2, f'Indexing on Points with {item} failed to return a matrix!'
        return original_type(p, points_dim=p.shape[1], attribute_dims=attribute_dims)

    def __len__(self):
        """int: Number of points in the current object."""
        return self.tensor.shape[0]

    def __repr__(self):
        """str: Return a strings that describes the object."""
        return self.__class__.__name__ + '(\n    ' + str(self.tensor) + ')'

    @classmethod
    def cat(cls, points_list):
        """Concatenate a list of Points into a single Points.

        Args:
            points_list (list[:obj:`BasePoints`]): List of points.

        Returns:
            :obj:`BasePoints`: The concatenated Points.
        """
        assert isinstance(points_list, (list, tuple))
        if len(points_list) == 0:
            return cls(torch.empty(0))
        assert all(isinstance(points, cls) for points in points_list)
        cat_points = cls(torch.cat([p.tensor for p in points_list], dim=0), points_dim=points_list[0].tensor.shape[1], attribute_dims=points_list[0].attribute_dims)
        return cat_points

    def to(self, device):
        """Convert current points to a specific device.

        Args:
            device (str | :obj:`torch.device`): The name of the device.

        Returns:
            :obj:`BasePoints`: A new boxes object on the
                specific device.
        """
        original_type = type(self)
        return original_type(self.tensor, points_dim=self.points_dim, attribute_dims=self.attribute_dims)

    def clone(self):
        """Clone the Points.

        Returns:
            :obj:`BasePoints`: Box object with the same properties
                as self.
        """
        original_type = type(self)
        return original_type(self.tensor.clone(), points_dim=self.points_dim, attribute_dims=self.attribute_dims)

    @property
    def device(self):
        """str: The device of the points are on."""
        return self.tensor.device

    def __iter__(self):
        """Yield a point as a Tensor of shape (4,) at a time.

        Returns:
            torch.Tensor: A point of shape (4,).
        """
        yield from self.tensor

    def new_point(self, data):
        """Create a new point object with data.

        The new point and its tensor has the similar properties
            as self and self.tensor, respectively.

        Args:
            data (torch.Tensor | numpy.array | list): Data to be copied.

        Returns:
            :obj:`BasePoints`: A new point object with ``data``,
                the object's other properties are similar to ``self``.
        """
        new_tensor = self.tensor.new_tensor(data) if not isinstance(data, torch.Tensor) else data
        original_type = type(self)
        return original_type(new_tensor, points_dim=self.points_dim, attribute_dims=self.attribute_dims)


class CameraPoints(BasePoints):
    """Points of instances in CAM coordinates.

    Args:
        tensor (torch.Tensor | np.ndarray | list): a N x points_dim matrix.
        points_dim (int, optional): Number of the dimension of a point.
            Each row is (x, y, z). Defaults to 3.
        attribute_dims (dict, optional): Dictionary to indicate the
            meaning of extra dimension. Defaults to None.

    Attributes:
        tensor (torch.Tensor): Float matrix of N x points_dim.
        points_dim (int): Integer indicating the dimension of a point.
            Each row is (x, y, z, ...).
        attribute_dims (bool): Dictionary to indicate the meaning of extra
            dimension. Defaults to None.
        rotation_axis (int): Default rotation axis for points rotation.
    """

    def __init__(self, tensor, points_dim=3, attribute_dims=None):
        super(CameraPoints, self).__init__(tensor, points_dim=points_dim, attribute_dims=attribute_dims)
        self.rotation_axis = 1

    def flip(self, bev_direction='horizontal'):
        """Flip the points along given BEV direction.

        Args:
            bev_direction (str): Flip direction (horizontal or vertical).
        """
        if bev_direction == 'horizontal':
            self.tensor[:, 0] = -self.tensor[:, 0]
        elif bev_direction == 'vertical':
            self.tensor[:, 2] = -self.tensor[:, 2]

    @property
    def bev(self):
        """torch.Tensor: BEV of the points in shape (N, 2)."""
        return self.tensor[:, [0, 2]]

    def convert_to(self, dst, rt_mat=None):
        """Convert self to ``dst`` mode.

        Args:
            dst (:obj:`CoordMode`): The target Point mode.
            rt_mat (np.ndarray | torch.Tensor, optional): The rotation and
                translation matrix between different coordinates.
                Defaults to None.
                The conversion from `src` coordinates to `dst` coordinates
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
            :obj:`BasePoints`: The converted point of the same type
                in the `dst` mode.
        """
        return Coord3DMode.convert_point(point=self, src=Coord3DMode.CAM, dst=dst, rt_mat=rt_mat)


class DepthPoints(BasePoints):
    """Points of instances in DEPTH coordinates.

    Args:
        tensor (torch.Tensor | np.ndarray | list): a N x points_dim matrix.
        points_dim (int, optional): Number of the dimension of a point.
            Each row is (x, y, z). Defaults to 3.
        attribute_dims (dict, optional): Dictionary to indicate the
            meaning of extra dimension. Defaults to None.

    Attributes:
        tensor (torch.Tensor): Float matrix of N x points_dim.
        points_dim (int): Integer indicating the dimension of a point.
            Each row is (x, y, z, ...).
        attribute_dims (bool): Dictionary to indicate the meaning of extra
            dimension. Defaults to None.
        rotation_axis (int): Default rotation axis for points rotation.
    """

    def __init__(self, tensor, points_dim=3, attribute_dims=None):
        super(DepthPoints, self).__init__(tensor, points_dim=points_dim, attribute_dims=attribute_dims)
        self.rotation_axis = 2

    def flip(self, bev_direction='horizontal'):
        """Flip the points along given BEV direction.

        Args:
            bev_direction (str): Flip direction (horizontal or vertical).
        """
        if bev_direction == 'horizontal':
            self.tensor[:, 0] = -self.tensor[:, 0]
        elif bev_direction == 'vertical':
            self.tensor[:, 1] = -self.tensor[:, 1]

    def convert_to(self, dst, rt_mat=None):
        """Convert self to ``dst`` mode.

        Args:
            dst (:obj:`CoordMode`): The target Point mode.
            rt_mat (np.ndarray | torch.Tensor, optional): The rotation and
                translation matrix between different coordinates.
                Defaults to None.
                The conversion from `src` coordinates to `dst` coordinates
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
            :obj:`BasePoints`: The converted point of the same type
                in the `dst` mode.
        """
        return Coord3DMode.convert_point(point=self, src=Coord3DMode.DEPTH, dst=dst, rt_mat=rt_mat)


class LiDARPoints(BasePoints):
    """Points of instances in LIDAR coordinates.

    Args:
        tensor (torch.Tensor | np.ndarray | list): a N x points_dim matrix.
        points_dim (int, optional): Number of the dimension of a point.
            Each row is (x, y, z). Defaults to 3.
        attribute_dims (dict, optional): Dictionary to indicate the
            meaning of extra dimension. Defaults to None.

    Attributes:
        tensor (torch.Tensor): Float matrix of N x points_dim.
        points_dim (int): Integer indicating the dimension of a point.
            Each row is (x, y, z, ...).
        attribute_dims (bool): Dictionary to indicate the meaning of extra
            dimension. Defaults to None.
        rotation_axis (int): Default rotation axis for points rotation.
    """

    def __init__(self, tensor, points_dim=3, attribute_dims=None):
        super(LiDARPoints, self).__init__(tensor, points_dim=points_dim, attribute_dims=attribute_dims)
        self.rotation_axis = 2

    def flip(self, bev_direction='horizontal'):
        """Flip the points along given BEV direction.

        Args:
            bev_direction (str): Flip direction (horizontal or vertical).
        """
        if bev_direction == 'horizontal':
            self.tensor[:, 1] = -self.tensor[:, 1]
        elif bev_direction == 'vertical':
            self.tensor[:, 0] = -self.tensor[:, 0]

    def convert_to(self, dst, rt_mat=None):
        """Convert self to ``dst`` mode.

        Args:
            dst (:obj:`CoordMode`): The target Point mode.
            rt_mat (np.ndarray | torch.Tensor, optional): The rotation and
                translation matrix between different coordinates.
                Defaults to None.
                The conversion from `src` coordinates to `dst` coordinates
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
            :obj:`BasePoints`: The converted point of the same type
                in the `dst` mode.
        """
        return Coord3DMode.convert_point(point=self, src=Coord3DMode.LIDAR, dst=dst, rt_mat=rt_mat)


def get_points_type(points_type):
    """Get the class of points according to coordinate type.

    Args:
        points_type (str): The type of points coordinate.
            The valid value are "CAMERA", "LIDAR", or "DEPTH".

    Returns:
        class: Points type.
    """
    if points_type == 'CAMERA':
        points_cls = CameraPoints
    elif points_type == 'LIDAR':
        points_cls = LiDARPoints
    elif points_type == 'DEPTH':
        points_cls = DepthPoints
    else:
        raise ValueError(f'Only "points_type" of "CAMERA", "LIDAR", or "DEPTH" are supported, got {points_type}')
    return points_cls


def apply_3d_transformation(pcd, coord_type, img_meta, reverse=False):
    """Apply transformation to input point cloud.

    Args:
        pcd (torch.Tensor): The point cloud to be transformed.
        coord_type (str): 'DEPTH' or 'CAMERA' or 'LIDAR'.
        img_meta(dict): Meta info regarding data transformation.
        reverse (bool): Reversed transformation or not.

    Note:
        The elements in img_meta['transformation_3d_flow']:
        "T" stands for translation;
        "S" stands for scale;
        "R" stands for rotation;
        "HF" stands for horizontal flip;
        "VF" stands for vertical flip.

    Returns:
        torch.Tensor: The transformed point cloud.
    """
    dtype = pcd.dtype
    device = pcd.device
    pcd_rotate_mat = torch.tensor(img_meta['pcd_rotation'], dtype=dtype, device=device) if 'pcd_rotation' in img_meta else torch.eye(3, dtype=dtype, device=device)
    pcd_scale_factor = img_meta['pcd_scale_factor'] if 'pcd_scale_factor' in img_meta else 1.0
    pcd_trans_factor = torch.tensor(img_meta['pcd_trans'], dtype=dtype, device=device) if 'pcd_trans' in img_meta else torch.zeros(3, dtype=dtype, device=device)
    pcd_horizontal_flip = img_meta['pcd_horizontal_flip'] if 'pcd_horizontal_flip' in img_meta else False
    pcd_vertical_flip = img_meta['pcd_vertical_flip'] if 'pcd_vertical_flip' in img_meta else False
    flow = img_meta['transformation_3d_flow'] if 'transformation_3d_flow' in img_meta else []
    pcd = pcd.clone()
    pcd = get_points_type(coord_type)(pcd)
    horizontal_flip_func = partial(pcd.flip, bev_direction='horizontal') if pcd_horizontal_flip else lambda : None
    vertical_flip_func = partial(pcd.flip, bev_direction='vertical') if pcd_vertical_flip else lambda : None
    if reverse:
        scale_func = partial(pcd.scale, scale_factor=1.0 / pcd_scale_factor)
        translate_func = partial(pcd.translate, trans_vector=-pcd_trans_factor)
        rotate_func = partial(pcd.rotate, rotation=pcd_rotate_mat.inverse())
        flow = flow[::-1]
    else:
        scale_func = partial(pcd.scale, scale_factor=pcd_scale_factor)
        translate_func = partial(pcd.translate, trans_vector=pcd_trans_factor)
        rotate_func = partial(pcd.rotate, rotation=pcd_rotate_mat)
    flow_mapping = {'T': translate_func, 'S': scale_func, 'R': rotate_func, 'HF': horizontal_flip_func, 'VF': vertical_flip_func}
    for op in flow:
        assert op in flow_mapping, f'This 3D data transformation op ({op}) is not supported'
        func = flow_mapping[op]
        func()
    return pcd.coord


def extract_2d_info(img_meta, tensor):
    """Extract image augmentation information from img_meta.

    Args:
        img_meta(dict): Meta info regarding data transformation.
        tensor(torch.Tensor): Input tensor used to create new ones.

    Returns:
        (int, int, int, int, torch.Tensor, bool, torch.Tensor):
            The extracted information.
    """
    img_shape = img_meta['img_shape']
    ori_shape = img_meta['ori_shape']
    img_h, img_w, _ = img_shape
    ori_h, ori_w, _ = ori_shape
    img_scale_factor = tensor.new_tensor(img_meta['scale_factor'][:2]) if 'scale_factor' in img_meta else tensor.new_tensor([1.0, 1.0])
    img_flip = img_meta['flip'] if 'flip' in img_meta else False
    img_crop_offset = tensor.new_tensor(img_meta['img_crop_offset']) if 'img_crop_offset' in img_meta else tensor.new_tensor([0.0, 0.0])
    return img_h, img_w, ori_h, ori_w, img_scale_factor, img_flip, img_crop_offset


def bbox_2d_transform(img_meta, bbox_2d, ori2new):
    """Transform 2d bbox according to img_meta.

    Args:
        img_meta(dict): Meta info regarding data transformation.
        bbox_2d (torch.Tensor): Shape (..., >4)
            The input 2d bboxes to transform.
        ori2new (bool): Origin img coord system to new or not.

    Returns:
        torch.Tensor: The transformed 2d bboxes.
    """
    img_h, img_w, ori_h, ori_w, img_scale_factor, img_flip, img_crop_offset = extract_2d_info(img_meta, bbox_2d)
    bbox_2d_new = bbox_2d.clone()
    if ori2new:
        bbox_2d_new[:, 0] = bbox_2d_new[:, 0] * img_scale_factor[0]
        bbox_2d_new[:, 2] = bbox_2d_new[:, 2] * img_scale_factor[0]
        bbox_2d_new[:, 1] = bbox_2d_new[:, 1] * img_scale_factor[1]
        bbox_2d_new[:, 3] = bbox_2d_new[:, 3] * img_scale_factor[1]
        bbox_2d_new[:, 0] = bbox_2d_new[:, 0] + img_crop_offset[0]
        bbox_2d_new[:, 2] = bbox_2d_new[:, 2] + img_crop_offset[0]
        bbox_2d_new[:, 1] = bbox_2d_new[:, 1] + img_crop_offset[1]
        bbox_2d_new[:, 3] = bbox_2d_new[:, 3] + img_crop_offset[1]
        if img_flip:
            bbox_2d_r = img_w - bbox_2d_new[:, 0]
            bbox_2d_l = img_w - bbox_2d_new[:, 2]
            bbox_2d_new[:, 0] = bbox_2d_l
            bbox_2d_new[:, 2] = bbox_2d_r
    else:
        if img_flip:
            bbox_2d_r = img_w - bbox_2d_new[:, 0]
            bbox_2d_l = img_w - bbox_2d_new[:, 2]
            bbox_2d_new[:, 0] = bbox_2d_l
            bbox_2d_new[:, 2] = bbox_2d_r
        bbox_2d_new[:, 0] = bbox_2d_new[:, 0] - img_crop_offset[0]
        bbox_2d_new[:, 2] = bbox_2d_new[:, 2] - img_crop_offset[0]
        bbox_2d_new[:, 1] = bbox_2d_new[:, 1] - img_crop_offset[1]
        bbox_2d_new[:, 3] = bbox_2d_new[:, 3] - img_crop_offset[1]
        bbox_2d_new[:, 0] = bbox_2d_new[:, 0] / img_scale_factor[0]
        bbox_2d_new[:, 2] = bbox_2d_new[:, 2] / img_scale_factor[0]
        bbox_2d_new[:, 1] = bbox_2d_new[:, 1] / img_scale_factor[1]
        bbox_2d_new[:, 3] = bbox_2d_new[:, 3] / img_scale_factor[1]
    return bbox_2d_new


def coord_2d_transform(img_meta, coord_2d, ori2new):
    """Transform 2d pixel coordinates according to img_meta.

    Args:
        img_meta(dict): Meta info regarding data transformation.
        coord_2d (torch.Tensor): Shape (..., 2)
            The input 2d coords to transform.
        ori2new (bool): Origin img coord system to new or not.

    Returns:
        torch.Tensor: The transformed 2d coordinates.
    """
    img_h, img_w, ori_h, ori_w, img_scale_factor, img_flip, img_crop_offset = extract_2d_info(img_meta, coord_2d)
    coord_2d_new = coord_2d.clone()
    if ori2new:
        coord_2d_new[..., 0] = coord_2d_new[..., 0] * img_scale_factor[0]
        coord_2d_new[..., 1] = coord_2d_new[..., 1] * img_scale_factor[1]
        coord_2d_new[..., 0] += img_crop_offset[0]
        coord_2d_new[..., 1] += img_crop_offset[1]
        if img_flip:
            coord_2d_new[..., 0] = img_w - coord_2d_new[..., 0]
    else:
        if img_flip:
            coord_2d_new[..., 0] = img_w - coord_2d_new[..., 0]
        coord_2d_new[..., 0] -= img_crop_offset[0]
        coord_2d_new[..., 1] -= img_crop_offset[1]
        coord_2d_new[..., 0] = coord_2d_new[..., 0] / img_scale_factor[0]
        coord_2d_new[..., 1] = coord_2d_new[..., 1] / img_scale_factor[1]
    return coord_2d_new


@array_converter(apply_to=('points_3d', 'proj_mat'))
def points_cam2img(points_3d, proj_mat, with_depth=False):
    """Project points in camera coordinates to image coordinates.

    Args:
        points_3d (torch.Tensor | np.ndarray): Points in shape (N, 3)
        proj_mat (torch.Tensor | np.ndarray):
            Transformation matrix between coordinates.
        with_depth (bool, optional): Whether to keep depth in the output.
            Defaults to False.

    Returns:
        (torch.Tensor | np.ndarray): Points in image coordinates,
            with shape [N, 2] if `with_depth=False`, else [N, 3].
    """
    points_shape = list(points_3d.shape)
    points_shape[-1] = 1
    assert len(proj_mat.shape) == 2, f'The dimension of the projection matrix should be 2 instead of {len(proj_mat.shape)}.'
    d1, d2 = proj_mat.shape[:2]
    assert d1 == 3 and d2 == 3 or d1 == 3 and d2 == 4 or d1 == 4 and d2 == 4, f'The shape of the projection matrix ({d1}*{d2}) is not supported.'
    if d1 == 3:
        proj_mat_expanded = torch.eye(4, device=proj_mat.device, dtype=proj_mat.dtype)
        proj_mat_expanded[:d1, :d2] = proj_mat
        proj_mat = proj_mat_expanded
    points_4 = torch.cat([points_3d, points_3d.new_ones(points_shape)], dim=-1)
    point_2d = points_4 @ proj_mat.T
    point_2d_res = point_2d[..., :2] / point_2d[..., 2:3]
    if with_depth:
        point_2d_res = torch.cat([point_2d_res, point_2d[..., 2:3]], dim=-1)
    return point_2d_res


class VoteFusion(nn.Module):
    """Fuse 2d features from 3d seeds.

    Args:
        num_classes (int): number of classes.
        max_imvote_per_pixel (int): max number of imvotes.
    """

    def __init__(self, num_classes=10, max_imvote_per_pixel=3):
        super(VoteFusion, self).__init__()
        self.num_classes = num_classes
        self.max_imvote_per_pixel = max_imvote_per_pixel

    def forward(self, imgs, bboxes_2d_rescaled, seeds_3d_depth, img_metas):
        """Forward function.

        Args:
            imgs (list[torch.Tensor]): Image features.
            bboxes_2d_rescaled (list[torch.Tensor]): 2D bboxes.
            seeds_3d_depth (torch.Tensor): 3D seeds.
            img_metas (list[dict]): Meta information of images.

        Returns:
            torch.Tensor: Concatenated cues of each point.
            torch.Tensor: Validity mask of each feature.
        """
        img_features = []
        masks = []
        for i, data in enumerate(zip(imgs, bboxes_2d_rescaled, seeds_3d_depth, img_metas)):
            img, bbox_2d_rescaled, seed_3d_depth, img_meta = data
            bbox_num = bbox_2d_rescaled.shape[0]
            seed_num = seed_3d_depth.shape[0]
            img_shape = img_meta['img_shape']
            img_h, img_w, _ = img_shape
            xyz_depth = apply_3d_transformation(seed_3d_depth, 'DEPTH', img_meta, reverse=True)
            depth2img = xyz_depth.new_tensor(img_meta['depth2img'])
            uvz_origin = points_cam2img(xyz_depth, depth2img, True)
            z_cam = uvz_origin[..., 2]
            uv_origin = (uvz_origin[..., :2] - 1).round()
            uv_rescaled = coord_2d_transform(img_meta, uv_origin, True)
            bbox_2d_origin = bbox_2d_transform(img_meta, bbox_2d_rescaled, False)
            if bbox_num == 0:
                imvote_num = seed_num * self.max_imvote_per_pixel
                two_cues = torch.zeros((15, imvote_num), device=seed_3d_depth.device)
                mask_zero = torch.zeros(imvote_num - seed_num, device=seed_3d_depth.device).bool()
                mask_one = torch.ones(seed_num, device=seed_3d_depth.device).bool()
                mask = torch.cat([mask_one, mask_zero], dim=0)
            else:
                bbox_expanded = bbox_2d_origin.view(1, bbox_num, -1).expand(seed_num, -1, -1)
                seed_2d_expanded = uv_origin.view(seed_num, 1, -1).expand(-1, bbox_num, -1)
                seed_2d_expanded_x, seed_2d_expanded_y = seed_2d_expanded.split(1, dim=-1)
                bbox_expanded_l, bbox_expanded_t, bbox_expanded_r, bbox_expanded_b, bbox_expanded_conf, bbox_expanded_cls = bbox_expanded.split(1, dim=-1)
                bbox_expanded_midx = (bbox_expanded_l + bbox_expanded_r) / 2
                bbox_expanded_midy = (bbox_expanded_t + bbox_expanded_b) / 2
                seed_2d_in_bbox_x = (seed_2d_expanded_x > bbox_expanded_l) * (seed_2d_expanded_x < bbox_expanded_r)
                seed_2d_in_bbox_y = (seed_2d_expanded_y > bbox_expanded_t) * (seed_2d_expanded_y < bbox_expanded_b)
                seed_2d_in_bbox = seed_2d_in_bbox_x * seed_2d_in_bbox_y
                sem_cue = torch.zeros_like(bbox_expanded_conf).expand(-1, -1, self.num_classes)
                sem_cue = sem_cue.scatter(-1, bbox_expanded_cls.long(), bbox_expanded_conf)
                delta_u = bbox_expanded_midx - seed_2d_expanded_x
                delta_v = bbox_expanded_midy - seed_2d_expanded_y
                seed_3d_expanded = seed_3d_depth.view(seed_num, 1, -1).expand(-1, bbox_num, -1)
                z_cam = z_cam.view(seed_num, 1, 1).expand(-1, bbox_num, -1)
                imvote = torch.cat([delta_u, delta_v, torch.zeros_like(delta_v)], dim=-1).view(-1, 3)
                imvote = imvote * z_cam.reshape(-1, 1)
                imvote = imvote @ torch.inverse(depth2img.t())
                imvote = apply_3d_transformation(imvote, 'DEPTH', img_meta, reverse=False)
                seed_3d_expanded = seed_3d_expanded.reshape(imvote.shape)
                ray_angle = seed_3d_expanded + imvote
                ray_angle /= torch.sqrt(torch.sum(ray_angle ** 2, -1) + EPS).unsqueeze(-1)
                xz = ray_angle[:, [0, 2]] / (ray_angle[:, [1]] + EPS) * seed_3d_expanded[:, [1]] - seed_3d_expanded[:, [0, 2]]
                geo_cue = torch.cat([xz, ray_angle], dim=-1).view(seed_num, -1, 5)
                two_cues = torch.cat([geo_cue, sem_cue], dim=-1)
                two_cues = two_cues * seed_2d_in_bbox.float()
                feature_size = two_cues.shape[-1]
                if bbox_num < self.max_imvote_per_pixel:
                    append_num = self.max_imvote_per_pixel - bbox_num
                    append_zeros = torch.zeros((seed_num, append_num, 1), device=seed_2d_in_bbox.device).bool()
                    seed_2d_in_bbox = torch.cat([seed_2d_in_bbox, append_zeros], dim=1)
                    append_zeros = torch.zeros((seed_num, append_num, feature_size), device=two_cues.device)
                    two_cues = torch.cat([two_cues, append_zeros], dim=1)
                    append_zeros = torch.zeros((seed_num, append_num, 1), device=two_cues.device)
                    bbox_expanded_conf = torch.cat([bbox_expanded_conf, append_zeros], dim=1)
                pair_score = seed_2d_in_bbox.float() + bbox_expanded_conf
                mask, indices = pair_score.topk(self.max_imvote_per_pixel, dim=1, largest=True, sorted=True)
                indices_img = indices.expand(-1, -1, feature_size)
                two_cues = two_cues.gather(dim=1, index=indices_img)
                two_cues = two_cues.transpose(1, 0)
                two_cues = two_cues.reshape(-1, feature_size).transpose(1, 0).contiguous()
                mask = mask.floor().int()
                mask = mask.transpose(1, 0).reshape(-1).bool()
            img = img[:, :img_shape[0], :img_shape[1]]
            img_flatten = img.reshape(3, -1).float()
            img_flatten /= 255.0
            uv_rescaled[:, 0] = torch.clamp(uv_rescaled[:, 0].round(), 0, img_shape[1] - 1)
            uv_rescaled[:, 1] = torch.clamp(uv_rescaled[:, 1].round(), 0, img_shape[0] - 1)
            uv_flatten = uv_rescaled[:, 1].round() * img_shape[1] + uv_rescaled[:, 0].round()
            uv_expanded = uv_flatten.unsqueeze(0).expand(3, -1).long()
            txt_cue = torch.gather(img_flatten, dim=-1, index=uv_expanded)
            txt_cue = txt_cue.unsqueeze(1).expand(-1, self.max_imvote_per_pixel, -1).reshape(3, -1)
            img_feature = torch.cat([two_cues, txt_cue], dim=0)
            img_features.append(img_feature)
            masks.append(mask)
        return torch.stack(img_features, 0), torch.stack(masks, 0)


def axis_aligned_bbox_overlaps_3d(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-06):
    """Calculate overlap between two set of axis aligned 3D bboxes. If
    ``is_aligned`` is ``False``, then calculate the overlaps between each bbox
    of bboxes1 and bboxes2, otherwise the overlaps between each aligned pair of
    bboxes1 and bboxes2.

    Args:
        bboxes1 (Tensor): shape (B, m, 6) in <x1, y1, z1, x2, y2, z2>
            format or empty.
        bboxes2 (Tensor): shape (B, n, 6) in <x1, y1, z1, x2, y2, z2>
            format or empty.
            B indicates the batch dim, in shape (B1, B2, ..., Bn).
            If ``is_aligned`` is ``True``, then m and n must be equal.
        mode (str): "iou" (intersection over union) or "giou" (generalized
            intersection over union).
        is_aligned (bool, optional): If True, then m and n must be equal.
            Defaults to False.
        eps (float, optional): A value added to the denominator for numerical
            stability. Defaults to 1e-6.

    Returns:
        Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)

    Example:
        >>> bboxes1 = torch.FloatTensor([
        >>>     [0, 0, 0, 10, 10, 10],
        >>>     [10, 10, 10, 20, 20, 20],
        >>>     [32, 32, 32, 38, 40, 42],
        >>> ])
        >>> bboxes2 = torch.FloatTensor([
        >>>     [0, 0, 0, 10, 20, 20],
        >>>     [0, 10, 10, 10, 19, 20],
        >>>     [10, 10, 10, 20, 20, 20],
        >>> ])
        >>> overlaps = axis_aligned_bbox_overlaps_3d(bboxes1, bboxes2)
        >>> assert overlaps.shape == (3, 3)
        >>> overlaps = bbox_overlaps(bboxes1, bboxes2, is_aligned=True)
        >>> assert overlaps.shape == (3, )
    Example:
        >>> empty = torch.empty(0, 6)
        >>> nonempty = torch.FloatTensor([[0, 0, 0, 10, 9, 10]])
        >>> assert tuple(bbox_overlaps(empty, nonempty).shape) == (0, 1)
        >>> assert tuple(bbox_overlaps(nonempty, empty).shape) == (1, 0)
        >>> assert tuple(bbox_overlaps(empty, empty).shape) == (0, 0)
    """
    assert mode in ['iou', 'giou'], f'Unsupported mode {mode}'
    assert bboxes1.size(-1) == 6 or bboxes1.size(0) == 0
    assert bboxes2.size(-1) == 6 or bboxes2.size(0) == 0
    assert bboxes1.shape[:-2] == bboxes2.shape[:-2]
    batch_shape = bboxes1.shape[:-2]
    rows = bboxes1.size(-2)
    cols = bboxes2.size(-2)
    if is_aligned:
        assert rows == cols
    if rows * cols == 0:
        if is_aligned:
            return bboxes1.new(batch_shape + (rows,))
        else:
            return bboxes1.new(batch_shape + (rows, cols))
    area1 = (bboxes1[..., 3] - bboxes1[..., 0]) * (bboxes1[..., 4] - bboxes1[..., 1]) * (bboxes1[..., 5] - bboxes1[..., 2])
    area2 = (bboxes2[..., 3] - bboxes2[..., 0]) * (bboxes2[..., 4] - bboxes2[..., 1]) * (bboxes2[..., 5] - bboxes2[..., 2])
    if is_aligned:
        lt = torch.max(bboxes1[..., :3], bboxes2[..., :3])
        rb = torch.min(bboxes1[..., 3:], bboxes2[..., 3:])
        wh = (rb - lt).clamp(min=0)
        overlap = wh[..., 0] * wh[..., 1] * wh[..., 2]
        if mode in ['iou', 'giou']:
            union = area1 + area2 - overlap
        else:
            union = area1
        if mode == 'giou':
            enclosed_lt = torch.min(bboxes1[..., :3], bboxes2[..., :3])
            enclosed_rb = torch.max(bboxes1[..., 3:], bboxes2[..., 3:])
    else:
        lt = torch.max(bboxes1[..., :, None, :3], bboxes2[..., None, :, :3])
        rb = torch.min(bboxes1[..., :, None, 3:], bboxes2[..., None, :, 3:])
        wh = (rb - lt).clamp(min=0)
        overlap = wh[..., 0] * wh[..., 1] * wh[..., 2]
        if mode in ['iou', 'giou']:
            union = area1[..., None] + area2[..., None, :] - overlap
        if mode == 'giou':
            enclosed_lt = torch.min(bboxes1[..., :, None, :3], bboxes2[..., None, :, :3])
            enclosed_rb = torch.max(bboxes1[..., :, None, 3:], bboxes2[..., None, :, 3:])
    eps = union.new_tensor([eps])
    union = torch.max(union, eps)
    ious = overlap / union
    if mode in ['iou']:
        return ious
    enclose_wh = (enclosed_rb - enclosed_lt).clamp(min=0)
    enclose_area = enclose_wh[..., 0] * enclose_wh[..., 1] * enclose_wh[..., 2]
    enclose_area = torch.max(enclose_area, eps)
    gious = ious - (enclose_area - union) / enclose_area
    return gious


class AxisAlignedBboxOverlaps3D(object):
    """Axis-aligned 3D Overlaps (IoU) Calculator."""

    def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False):
        """Calculate IoU between 2D bboxes.

        Args:
            bboxes1 (Tensor): shape (B, m, 6) in <x1, y1, z1, x2, y2, z2>
                format or empty.
            bboxes2 (Tensor): shape (B, n, 6) in <x1, y1, z1, x2, y2, z2>
                format or empty.
                B indicates the batch dim, in shape (B1, B2, ..., Bn).
                If ``is_aligned`` is ``True``, then m and n must be equal.
            mode (str): "iou" (intersection over union) or "giou" (generalized
                intersection over union).
            is_aligned (bool, optional): If True, then m and n must be equal.
                Defaults to False.
        Returns:
            Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)
        """
        assert bboxes1.size(-1) == bboxes2.size(-1) == 6
        return axis_aligned_bbox_overlaps_3d(bboxes1, bboxes2, mode, is_aligned)

    def __repr__(self):
        """str: a string describing the module"""
        repr_str = self.__class__.__name__ + '()'
        return repr_str


class AxisAlignedIoULoss(nn.Module):
    """Calculate the IoU loss (1-IoU) of axis aligned bounding boxes.

    Args:
        reduction (str): Method to reduce losses.
            The valid reduction method are none, sum or mean.
        loss_weight (float, optional): Weight of loss. Defaults to 1.0.
    """

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(AxisAlignedIoULoss, self).__init__()
        assert reduction in ['none', 'sum', 'mean']
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, weight=None, avg_factor=None, reduction_override=None, **kwargs):
        """Forward function of loss calculation.

        Args:
            pred (torch.Tensor): Bbox predictions with shape [..., 6]
                (x1, y1, z1, x2, y2, z2).
            target (torch.Tensor): Bbox targets (gt) with shape [..., 6]
                (x1, y1, z1, x2, y2, z2).
            weight (torch.Tensor | float, optional): Weight of loss.
                Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): Method to reduce losses.
                The valid reduction method are 'none', 'sum' or 'mean'.
                Defaults to None.

        Returns:
            torch.Tensor: IoU loss between predictions and targets.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        if weight is not None and not torch.any(weight > 0) and reduction != 'none':
            return (pred * weight).sum()
        return axis_aligned_iou_loss(pred, target, weight=weight, avg_factor=avg_factor, reduction=reduction) * self.loss_weight


def chamfer_distance(src, dst, src_weight=1.0, dst_weight=1.0, criterion_mode='l2', reduction='mean'):
    """Calculate Chamfer Distance of two sets.

    Args:
        src (torch.Tensor): Source set with shape [B, N, C] to
            calculate Chamfer Distance.
        dst (torch.Tensor): Destination set with shape [B, M, C] to
            calculate Chamfer Distance.
        src_weight (torch.Tensor or float): Weight of source loss.
        dst_weight (torch.Tensor or float): Weight of destination loss.
        criterion_mode (str): Criterion mode to calculate distance.
            The valid modes are smooth_l1, l1 or l2.
        reduction (str): Method to reduce losses.
            The valid reduction method are 'none', 'sum' or 'mean'.

    Returns:
        tuple: Source and Destination loss with the corresponding indices.

            - loss_src (torch.Tensor): The min distance
                from source to destination.
            - loss_dst (torch.Tensor): The min distance
                from destination to source.
            - indices1 (torch.Tensor): Index the min distance point
                for each point in source to destination.
            - indices2 (torch.Tensor): Index the min distance point
                for each point in destination to source.
    """
    if criterion_mode == 'smooth_l1':
        criterion = smooth_l1_loss
    elif criterion_mode == 'l1':
        criterion = l1_loss
    elif criterion_mode == 'l2':
        criterion = mse_loss
    else:
        raise NotImplementedError
    src_expand = src.unsqueeze(2).repeat(1, 1, dst.shape[1], 1)
    dst_expand = dst.unsqueeze(1).repeat(1, src.shape[1], 1, 1)
    distance = criterion(src_expand, dst_expand, reduction='none').sum(-1)
    src2dst_distance, indices1 = torch.min(distance, dim=2)
    dst2src_distance, indices2 = torch.min(distance, dim=1)
    loss_src = src2dst_distance * src_weight
    loss_dst = dst2src_distance * dst_weight
    if reduction == 'sum':
        loss_src = torch.sum(loss_src)
        loss_dst = torch.sum(loss_dst)
    elif reduction == 'mean':
        loss_src = torch.mean(loss_src)
        loss_dst = torch.mean(loss_dst)
    elif reduction == 'none':
        pass
    else:
        raise NotImplementedError
    return loss_src, loss_dst, indices1, indices2


class ChamferDistance(nn.Module):
    """Calculate Chamfer Distance of two sets.

    Args:
        mode (str): Criterion mode to calculate distance.
            The valid modes are smooth_l1, l1 or l2.
        reduction (str): Method to reduce losses.
            The valid reduction method are none, sum or mean.
        loss_src_weight (float): Weight of loss_source.
        loss_dst_weight (float): Weight of loss_target.
    """

    def __init__(self, mode='l2', reduction='mean', loss_src_weight=1.0, loss_dst_weight=1.0):
        super(ChamferDistance, self).__init__()
        assert mode in ['smooth_l1', 'l1', 'l2']
        assert reduction in ['none', 'sum', 'mean']
        self.mode = mode
        self.reduction = reduction
        self.loss_src_weight = loss_src_weight
        self.loss_dst_weight = loss_dst_weight

    def forward(self, source, target, src_weight=1.0, dst_weight=1.0, reduction_override=None, return_indices=False, **kwargs):
        """Forward function of loss calculation.

        Args:
            source (torch.Tensor): Source set with shape [B, N, C] to
                calculate Chamfer Distance.
            target (torch.Tensor): Destination set with shape [B, M, C] to
                calculate Chamfer Distance.
            src_weight (torch.Tensor | float, optional):
                Weight of source loss. Defaults to 1.0.
            dst_weight (torch.Tensor | float, optional):
                Weight of destination loss. Defaults to 1.0.
            reduction_override (str, optional): Method to reduce losses.
                The valid reduction method are 'none', 'sum' or 'mean'.
                Defaults to None.
            return_indices (bool, optional): Whether to return indices.
                Defaults to False.

        Returns:
            tuple[torch.Tensor]: If ``return_indices=True``, return losses of
                source and target with their corresponding indices in the
                order of ``(loss_source, loss_target, indices1, indices2)``.
                If ``return_indices=False``, return
                ``(loss_source, loss_target)``.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        loss_source, loss_target, indices1, indices2 = chamfer_distance(source, target, src_weight, dst_weight, self.mode, reduction)
        loss_source *= self.loss_src_weight
        loss_target *= self.loss_dst_weight
        if return_indices:
            return loss_source, loss_target, indices1, indices2
        else:
            return loss_source, loss_target


class MultiBinLoss(nn.Module):
    """Multi-Bin Loss for orientation.

    Args:
        reduction (str, optional): The method to reduce the loss.
            Options are 'none', 'mean' and 'sum'. Defaults to 'none'.
        loss_weight (float, optional): The weight of loss. Defaults
            to 1.0.
    """

    def __init__(self, reduction='none', loss_weight=1.0):
        super(MultiBinLoss, self).__init__()
        assert reduction in ['none', 'sum', 'mean']
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, num_dir_bins, reduction_override=None):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            num_dir_bins (int): Number of bins to encode direction angle.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        loss = self.loss_weight * multibin_loss(pred, target, num_dir_bins=num_dir_bins, reduction=reduction)
        return loss


class ScoreNet(nn.Module):
    """ScoreNet that outputs coefficient scores to assemble kernel weights in
    the weight bank according to the relative position of point pairs.

    Args:
        mlp_channels (List[int]): Hidden unit sizes of SharedMLP layers.
        last_bn (bool, optional): Whether to use BN on the last output of mlps.
            Defaults to False.
        score_norm (str, optional): Normalization function of output scores.
            Can be 'softmax', 'sigmoid' or 'identity'. Defaults to 'softmax'.
        temp_factor (float, optional): Temperature factor to scale the output
            scores before softmax. Defaults to 1.0.
        norm_cfg (dict, optional): Type of normalization method.
            Defaults to dict(type='BN2d').
        bias (bool | str, optional): If specified as `auto`, it will be decided
            by the norm_cfg. Bias will be set as True if `norm_cfg` is None,
            otherwise False. Defaults to 'auto'.

    Note:
        The official code applies xavier_init to all Conv layers in ScoreNet,
            see `PAConv <https://github.com/CVMI-Lab/PAConv/blob/main/scene_seg
            /model/pointnet2/paconv.py#L105>`_. However in our experiments, we
            did not find much difference in applying such xavier initialization
            or not. So we neglect this initialization in our implementation.
    """

    def __init__(self, mlp_channels, last_bn=False, score_norm='softmax', temp_factor=1.0, norm_cfg=dict(type='BN2d'), bias='auto'):
        super(ScoreNet, self).__init__()
        assert score_norm in ['softmax', 'sigmoid', 'identity'], f'unsupported score_norm function {score_norm}'
        self.score_norm = score_norm
        self.temp_factor = temp_factor
        self.mlps = nn.Sequential()
        for i in range(len(mlp_channels) - 2):
            self.mlps.add_module(f'layer{i}', ConvModule(mlp_channels[i], mlp_channels[i + 1], kernel_size=(1, 1), stride=(1, 1), conv_cfg=dict(type='Conv2d'), norm_cfg=norm_cfg, bias=bias))
        i = len(mlp_channels) - 2
        self.mlps.add_module(f'layer{i}', ConvModule(mlp_channels[i], mlp_channels[i + 1], kernel_size=(1, 1), stride=(1, 1), conv_cfg=dict(type='Conv2d'), norm_cfg=norm_cfg if last_bn else None, act_cfg=None, bias=bias))

    def forward(self, xyz_features):
        """Forward.

        Args:
            xyz_features (torch.Tensor): (B, C, N, K), features constructed
                from xyz coordinates of point pairs. May contain relative
                positions, Euclidean distance, etc.

        Returns:
            torch.Tensor: (B, N, K, M), predicted scores for `M` kernels.
        """
        scores = self.mlps(xyz_features)
        if self.score_norm == 'softmax':
            scores = F.softmax(scores / self.temp_factor, dim=1)
        elif self.score_norm == 'sigmoid':
            scores = torch.sigmoid(scores / self.temp_factor)
        else:
            scores = scores
        scores = scores.permute(0, 2, 3, 1)
        return scores


def assign_score(scores, point_features):
    """Perform weighted sum to aggregate output features according to scores.
    This function is used in non-CUDA version of PAConv.

    Compared to the cuda op assigh_score_withk, this pytorch implementation
        pre-computes output features for the neighbors of all centers, and then
        performs aggregation. It consumes more GPU memories.

    Args:
        scores (torch.Tensor): (B, npoint, K, M), predicted scores to
            aggregate weight matrices in the weight bank.
            `npoint` is the number of sampled centers.
            `K` is the number of queried neighbors.
            `M` is the number of weight matrices in the weight bank.
        point_features (torch.Tensor): (B, npoint, K, M, out_dim)
            Pre-computed point features to be aggregated.

    Returns:
        torch.Tensor: (B, npoint, K, out_dim), the aggregated features.
    """
    B, npoint, K, M = scores.size()
    scores = scores.view(B, npoint, K, 1, M)
    output = torch.matmul(scores, point_features).view(B, npoint, K, -1)
    return output


def calc_euclidian_dist(xyz1, xyz2):
    """Calculate the Euclidean distance between two sets of points.

    Args:
        xyz1 (torch.Tensor): (N, 3), the first set of points.
        xyz2 (torch.Tensor): (N, 3), the second set of points.

    Returns:
        torch.Tensor: (N, ), the Euclidean distance between each point pair.
    """
    assert xyz1.shape[0] == xyz2.shape[0], 'number of points are not the same'
    assert xyz1.shape[1] == xyz2.shape[1] == 3, 'points coordinates dimension is not 3'
    return torch.norm(xyz1 - xyz2, dim=-1)


class PAConv(nn.Module):
    """Non-CUDA version of PAConv.

    PAConv stores a trainable weight bank containing several kernel weights.
    Given input points and features, it computes coefficient scores to assemble
    those kernels to form conv kernels, and then runs convolution on the input.

    Args:
        in_channels (int): Input channels of point features.
        out_channels (int): Output channels of point features.
        num_kernels (int): Number of kernel weights in the weight bank.
        norm_cfg (dict, optional): Type of normalization method.
            Defaults to dict(type='BN2d', momentum=0.1).
        act_cfg (dict, optional): Type of activation method.
            Defaults to dict(type='ReLU', inplace=True).
        scorenet_input (str, optional): Type of input to ScoreNet.
            Can be 'identity', 'w_neighbor' or 'w_neighbor_dist'.
            Defaults to 'w_neighbor_dist'.
        weight_bank_init (str, optional): Init method of weight bank kernels.
            Can be 'kaiming' or 'xavier'. Defaults to 'kaiming'.
        kernel_input (str, optional): Input features to be multiplied with
            kernel weights. Can be 'identity' or 'w_neighbor'.
            Defaults to 'w_neighbor'.
        scorenet_cfg (dict, optional): Config of the ScoreNet module, which
            may contain the following keys and values:

            - mlp_channels (List[int]): Hidden units of MLPs.
            - score_norm (str): Normalization function of output scores.
                Can be 'softmax', 'sigmoid' or 'identity'.
            - temp_factor (float): Temperature factor to scale the output
                scores before softmax.
            - last_bn (bool): Whether to use BN on the last output of mlps.
    """

    def __init__(self, in_channels, out_channels, num_kernels, norm_cfg=dict(type='BN2d', momentum=0.1), act_cfg=dict(type='ReLU', inplace=True), scorenet_input='w_neighbor_dist', weight_bank_init='kaiming', kernel_input='w_neighbor', scorenet_cfg=dict(mlp_channels=[16, 16, 16], score_norm='softmax', temp_factor=1.0, last_bn=False)):
        super(PAConv, self).__init__()
        if kernel_input == 'identity':
            kernel_mul = 1
        elif kernel_input == 'w_neighbor':
            kernel_mul = 2
        else:
            raise NotImplementedError(f'unsupported kernel_input {kernel_input}')
        self.kernel_input = kernel_input
        in_channels = kernel_mul * in_channels
        if scorenet_input == 'identity':
            self.scorenet_in_channels = 3
        elif scorenet_input == 'w_neighbor':
            self.scorenet_in_channels = 6
        elif scorenet_input == 'w_neighbor_dist':
            self.scorenet_in_channels = 7
        else:
            raise NotImplementedError(f'unsupported scorenet_input {scorenet_input}')
        self.scorenet_input = scorenet_input
        if weight_bank_init == 'kaiming':
            weight_init = nn.init.kaiming_normal_
        elif weight_bank_init == 'xavier':
            weight_init = nn.init.xavier_normal_
        else:
            raise NotImplementedError(f'unsupported weight bank init method {weight_bank_init}')
        self.num_kernels = num_kernels
        weight_bank = weight_init(torch.empty(self.num_kernels, in_channels, out_channels))
        weight_bank = weight_bank.permute(1, 0, 2).reshape(in_channels, self.num_kernels * out_channels).contiguous()
        self.weight_bank = nn.Parameter(weight_bank, requires_grad=True)
        scorenet_cfg_ = copy.deepcopy(scorenet_cfg)
        scorenet_cfg_['mlp_channels'].insert(0, self.scorenet_in_channels)
        scorenet_cfg_['mlp_channels'].append(self.num_kernels)
        self.scorenet = ScoreNet(**scorenet_cfg_)
        self.bn = build_norm_layer(norm_cfg, out_channels)[1] if norm_cfg is not None else None
        self.activate = build_activation_layer(act_cfg) if act_cfg is not None else None
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.init_weights()

    def init_weights(self):
        """Initialize weights of shared MLP layers and BN layers."""
        if self.bn is not None:
            constant_init(self.bn, val=1, bias=0)

    def _prepare_scorenet_input(self, points_xyz):
        """Prepare input point pairs features for self.ScoreNet.

        Args:
            points_xyz (torch.Tensor): (B, 3, npoint, K)
                Coordinates of the grouped points.

        Returns:
            torch.Tensor: (B, C, npoint, K)
                The generated features per point pair.
        """
        B, _, npoint, K = points_xyz.size()
        center_xyz = points_xyz[..., :1].repeat(1, 1, 1, K)
        xyz_diff = points_xyz - center_xyz
        if self.scorenet_input == 'identity':
            xyz_features = xyz_diff
        elif self.scorenet_input == 'w_neighbor':
            xyz_features = torch.cat((xyz_diff, points_xyz), dim=1)
        else:
            euclidian_dist = calc_euclidian_dist(center_xyz.permute(0, 2, 3, 1).reshape(B * npoint * K, 3), points_xyz.permute(0, 2, 3, 1).reshape(B * npoint * K, 3)).reshape(B, 1, npoint, K)
            xyz_features = torch.cat((center_xyz, xyz_diff, euclidian_dist), dim=1)
        return xyz_features

    def forward(self, inputs):
        """Forward.

        Args:
            inputs (tuple(torch.Tensor)):

                - features (torch.Tensor): (B, in_c, npoint, K)
                    Features of the queried points.
                - points_xyz (torch.Tensor): (B, 3, npoint, K)
                    Coordinates of the grouped points.

        Returns:
            Tuple[torch.Tensor]:

                - new_features: (B, out_c, npoint, K), features after PAConv.
                - points_xyz: same as input.
        """
        features, points_xyz = inputs
        B, _, npoint, K = features.size()
        if self.kernel_input == 'w_neighbor':
            center_features = features[..., :1].repeat(1, 1, 1, K)
            features_diff = features - center_features
            features = torch.cat((features_diff, features), dim=1)
        xyz_features = self._prepare_scorenet_input(points_xyz)
        scores = self.scorenet(xyz_features)
        new_features = torch.matmul(features.permute(0, 2, 3, 1), self.weight_bank).view(B, npoint, K, self.num_kernels, -1)
        new_features = assign_score(scores, new_features)
        new_features = new_features.permute(0, 3, 1, 2).contiguous()
        if self.bn is not None:
            new_features = self.bn(new_features)
        if self.activate is not None:
            new_features = self.activate(new_features)
        return new_features, points_xyz


def assign_kernel_withoutk(features, kernels, M):
    """Pre-compute features with weight matrices in weight bank. This function
    is used before cuda op assign_score_withk in CUDA version PAConv.

    Args:
        features (torch.Tensor): (B, in_dim, N), input features of all points.
            `N` is the number of points in current point cloud.
        kernels (torch.Tensor): (2 * in_dim, M * out_dim), weight matrices in
            the weight bank, transformed from (M, 2 * in_dim, out_dim).
            `2 * in_dim` is because the input features are concatenation of
            (point_features - center_features, point_features).
        M (int): Number of weight matrices in the weight bank.

    Returns:
        Tuple[torch.Tensor]: both of shape (B, N, M, out_dim):

            - point_features: Pre-computed features for points.
            - center_features: Pre-computed features for centers.
    """
    B, in_dim, N = features.size()
    feat_trans = features.permute(0, 2, 1)
    out_feat_half1 = torch.matmul(feat_trans, kernels[:in_dim]).view(B, N, M, -1)
    out_feat_half2 = torch.matmul(feat_trans, kernels[in_dim:]).view(B, N, M, -1)
    if features.size(1) % 2 != 0:
        out_feat_half_coord = torch.matmul(feat_trans[:, :, :3], kernels[in_dim:in_dim + 3]).view(B, N, M, -1)
    else:
        out_feat_half_coord = torch.zeros_like(out_feat_half2)
    point_features = out_feat_half1 + out_feat_half2
    center_features = out_feat_half1 + out_feat_half_coord
    return point_features, center_features


class PAConvCUDA(PAConv):
    """CUDA version of PAConv that implements a cuda op to efficiently perform
    kernel assembling.

    Different from vanilla PAConv, the input features of this function is not
    grouped by centers. Instead, they will be queried on-the-fly by the
    additional input `points_idx`. This avoids the large intermediate matrix.
    See the `paper <https://arxiv.org/pdf/2103.14635.pdf>`_ appendix Sec. D for
    more detailed descriptions.
    """

    def __init__(self, in_channels, out_channels, num_kernels, norm_cfg=dict(type='BN2d', momentum=0.1), act_cfg=dict(type='ReLU', inplace=True), scorenet_input='w_neighbor_dist', weight_bank_init='kaiming', kernel_input='w_neighbor', scorenet_cfg=dict(mlp_channels=[8, 16, 16], score_norm='softmax', temp_factor=1.0, last_bn=False)):
        super(PAConvCUDA, self).__init__(in_channels=in_channels, out_channels=out_channels, num_kernels=num_kernels, norm_cfg=norm_cfg, act_cfg=act_cfg, scorenet_input=scorenet_input, weight_bank_init=weight_bank_init, kernel_input=kernel_input, scorenet_cfg=scorenet_cfg)
        assert self.kernel_input == 'w_neighbor', 'CUDA implemented PAConv only supports w_neighbor kernel_input'

    def forward(self, inputs):
        """Forward.

        Args:
            inputs (tuple(torch.Tensor)):

                - features (torch.Tensor): (B, in_c, N)
                    Features of all points in the current point cloud.
                    Different from non-CUDA version PAConv, here the features
                        are not grouped by each center to form a K dim.
                - points_xyz (torch.Tensor): (B, 3, npoint, K)
                    Coordinates of the grouped points.
                - points_idx (torch.Tensor): (B, npoint, K)
                    Index of the grouped points.

        Returns:
            Tuple[torch.Tensor]:

                - new_features: (B, out_c, npoint, K), features after PAConv.
                - points_xyz: same as input.
                - points_idx: same as input.
        """
        features, points_xyz, points_idx = inputs
        xyz_features = self._prepare_scorenet_input(points_xyz)
        scores = self.scorenet(xyz_features)
        point_feat, center_feat = assign_kernel_withoutk(features, self.weight_bank, self.num_kernels)
        new_features = assign_score_cuda(scores, point_feat, center_feat, points_idx, 'sum').contiguous()
        if self.bn is not None:
            new_features = self.bn(new_features)
        if self.activate is not None:
            new_features = self.activate(new_features)
        return new_features, points_xyz, points_idx


def weight_correlation(conv):
    """Calculate correlations between kernel weights in Conv's weight bank as
    regularization loss. The cosine similarity is used as metrics.

    Args:
        conv (nn.Module): A Conv modules to be regularized.
            Currently we only support `PAConv` and `PAConvCUDA`.

    Returns:
        torch.Tensor: Correlations between each kernel weights in weight bank.
    """
    assert isinstance(conv, (PAConv, PAConvCUDA)), f'unsupported module type {type(conv)}'
    kernels = conv.weight_bank
    in_channels = conv.in_channels
    out_channels = conv.out_channels
    num_kernels = conv.num_kernels
    flatten_kernels = kernels.view(in_channels, num_kernels, out_channels).permute(1, 0, 2).reshape(num_kernels, -1)
    inner_product = torch.matmul(flatten_kernels, flatten_kernels.T)
    kernel_norms = torch.sum(flatten_kernels ** 2, dim=-1, keepdim=True) ** 0.5
    kernel_norms = torch.matmul(kernel_norms, kernel_norms.T)
    cosine_sims = inner_product / kernel_norms
    corr = torch.sum(torch.triu(cosine_sims, diagonal=1) ** 2)
    return corr


def paconv_regularization_loss(modules, reduction):
    """Computes correlation loss of PAConv weight kernels as regularization.

    Args:
        modules (List[nn.Module] | :obj:`generator`):
            A list or a python generator of torch.nn.Modules.
        reduction (str): Method to reduce losses among PAConv modules.
            The valid reduction method are none, sum or mean.

    Returns:
        torch.Tensor: Correlation loss of kernel weights.
    """
    corr_loss = []
    for module in modules:
        if isinstance(module, (PAConv, PAConvCUDA)):
            corr_loss.append(weight_correlation(module))
    corr_loss = torch.stack(corr_loss)
    corr_loss = weight_reduce_loss(corr_loss, reduction=reduction)
    return corr_loss


class PAConvRegularizationLoss(nn.Module):
    """Calculate correlation loss of kernel weights in PAConv's weight bank.

    This is used as a regularization term in PAConv model training.

    Args:
        reduction (str): Method to reduce losses. The reduction is performed
            among all PAConv modules instead of prediction tensors.
            The valid reduction method are none, sum or mean.
        loss_weight (float, optional): Weight of loss. Defaults to 1.0.
    """

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(PAConvRegularizationLoss, self).__init__()
        assert reduction in ['none', 'sum', 'mean']
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, modules, reduction_override=None, **kwargs):
        """Forward function of loss calculation.

        Args:
            modules (List[nn.Module] | :obj:`generator`):
                A list or a python generator of torch.nn.Modules.
            reduction_override (str, optional): Method to reduce losses.
                The valid reduction method are 'none', 'sum' or 'mean'.
                Defaults to None.

        Returns:
            torch.Tensor: Correlation loss of kernel weights.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        return self.loss_weight * paconv_regularization_loss(modules, reduction=reduction)


class RotatedIoU3DLoss(nn.Module):
    """Calculate the IoU loss (1-IoU) of rotated bounding boxes.

    Args:
        reduction (str): Method to reduce losses.
            The valid reduction method are none, sum or mean.
        loss_weight (float, optional): Weight of loss. Defaults to 1.0.
    """

    def __init__(self, reduction='mean', loss_weight=1.0):
        super().__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, weight=None, avg_factor=None, reduction_override=None, **kwargs):
        """Forward function of loss calculation.

        Args:
            pred (torch.Tensor): Bbox predictions with shape [..., 7]
                (x, y, z, w, l, h, alpha).
            target (torch.Tensor): Bbox targets (gt) with shape [..., 7]
                (x, y, z, w, l, h, alpha).
            weight (torch.Tensor | float, optional): Weight of loss.
                Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): Method to reduce losses.
                The valid reduction method are 'none', 'sum' or 'mean'.
                Defaults to None.

        Returns:
            torch.Tensor: IoU loss between predictions and targets.
        """
        if weight is not None and not torch.any(weight > 0):
            return pred.sum() * weight.sum()
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        if weight is not None and weight.dim() > 1:
            weight = weight.mean(-1)
        loss = self.loss_weight * rotated_iou_3d_loss(pred, target, weight, reduction=reduction, avg_factor=avg_factor, **kwargs)
        return loss


class UncertainSmoothL1Loss(nn.Module):
    """Smooth L1 loss with uncertainty.

    Please refer to `PGD <https://arxiv.org/abs/2107.14160>`_ and
    `Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry
    and Semantics <https://arxiv.org/abs/1705.07115>`_ for more details.

    Args:
        alpha (float, optional): The coefficient of log(sigma).
            Defaults to 1.0.
        beta (float, optional): The threshold in the piecewise function.
            Defaults to 1.0.
        reduction (str, optional): The method to reduce the loss.
            Options are 'none', 'mean' and 'sum'. Defaults to 'mean'.
        loss_weight (float, optional): The weight of loss. Defaults to 1.0
    """

    def __init__(self, alpha=1.0, beta=1.0, reduction='mean', loss_weight=1.0):
        super(UncertainSmoothL1Loss, self).__init__()
        assert reduction in ['none', 'sum', 'mean']
        self.alpha = alpha
        self.beta = beta
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, sigma, weight=None, avg_factor=None, reduction_override=None, **kwargs):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            sigma (torch.Tensor): The sigma for uncertainty.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        loss_bbox = self.loss_weight * uncertain_smooth_l1_loss(pred, target, weight, sigma=sigma, alpha=self.alpha, beta=self.beta, reduction=reduction, avg_factor=avg_factor, **kwargs)
        return loss_bbox


class UncertainL1Loss(nn.Module):
    """L1 loss with uncertainty.

    Args:
        alpha (float, optional): The coefficient of log(sigma).
            Defaults to 1.0.
        reduction (str, optional): The method to reduce the loss.
            Options are 'none', 'mean' and 'sum'. Defaults to 'mean'.
        loss_weight (float, optional): The weight of loss. Defaults to 1.0.
    """

    def __init__(self, alpha=1.0, reduction='mean', loss_weight=1.0):
        super(UncertainL1Loss, self).__init__()
        assert reduction in ['none', 'sum', 'mean']
        self.alpha = alpha
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self, pred, target, sigma, weight=None, avg_factor=None, reduction_override=None):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            sigma (torch.Tensor): The sigma for uncertainty.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = reduction_override if reduction_override else self.reduction
        loss_bbox = self.loss_weight * uncertain_l1_loss(pred, target, weight, sigma=sigma, alpha=self.alpha, reduction=reduction, avg_factor=avg_factor)
        return loss_bbox


def replace_feature(out, new_features):
    if 'replace_feature' in out.__dir__():
        return out.replace_feature(new_features)
    else:
        out.features = new_features
        return out


def make_sparse_convmodule(in_channels, out_channels, kernel_size, indice_key, stride=1, padding=0, conv_type='SubMConv3d', norm_cfg=None, order=('conv', 'norm', 'act')):
    """Make sparse convolution module.

    Args:
        in_channels (int): the number of input channels
        out_channels (int): the number of out channels
        kernel_size (int|tuple(int)): kernel size of convolution
        indice_key (str): the indice key used for sparse tensor
        stride (int|tuple(int)): the stride of convolution
        padding (int or list[int]): the padding number of input
        conv_type (str): sparse conv type in spconv
        norm_cfg (dict[str]): config of normalization layer
        order (tuple[str]): The order of conv/norm/activation layers. It is a
            sequence of "conv", "norm" and "act". Common examples are
            ("conv", "norm", "act") and ("act", "conv", "norm").

    Returns:
        spconv.SparseSequential: sparse convolution module.
    """
    assert isinstance(order, tuple) and len(order) <= 3
    assert set(order) | {'conv', 'norm', 'act'} == {'conv', 'norm', 'act'}
    conv_cfg = dict(type=conv_type, indice_key=indice_key)
    layers = list()
    for layer in order:
        if layer == 'conv':
            if conv_type not in ['SparseInverseConv3d', 'SparseInverseConv2d', 'SparseInverseConv1d']:
                layers.append(build_conv_layer(conv_cfg, in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=False))
            else:
                layers.append(build_conv_layer(conv_cfg, in_channels, out_channels, kernel_size, bias=False))
        elif layer == 'norm':
            layers.append(build_norm_layer(norm_cfg, out_channels)[1])
        elif layer == 'act':
            layers.append(nn.ReLU(inplace=True))
    layers = SparseSequential(*layers)
    return layers


class ConvBNPositionalEncoding(nn.Module):
    """Absolute position embedding with Conv learning.

    Args:
        input_channel (int): input features dim.
        num_pos_feats (int, optional): output position features dim.
            Defaults to 288 to be consistent with seed features dim.
    """

    def __init__(self, input_channel, num_pos_feats=288):
        super().__init__()
        self.position_embedding_head = nn.Sequential(nn.Conv1d(input_channel, num_pos_feats, kernel_size=1), nn.BatchNorm1d(num_pos_feats), nn.ReLU(inplace=True), nn.Conv1d(num_pos_feats, num_pos_feats, kernel_size=1))

    def forward(self, xyz):
        """Forward pass.

        Args:
            xyz (Tensor)： (B, N, 3) the coordinates to embed.

        Returns:
            Tensor: (B, num_pos_feats, N) the embedded position features.
        """
        xyz = xyz.permute(0, 2, 1)
        position_embedding = self.position_embedding_head(xyz)
        return position_embedding


def build_loss(cfg):
    """Build loss function."""
    if cfg['type'] in LOSSES._module_dict.keys():
        return LOSSES.build(cfg)
    elif cfg['type'] in MMDET_LOSSES._module_dict.keys():
        return MMDET_LOSSES.build(cfg)
    else:
        return MMSEG_LOSSES.build(cfg)


class VoteModule(nn.Module):
    """Vote module.

    Generate votes from seed point features.

    Args:
        in_channels (int): Number of channels of seed point features.
        vote_per_seed (int, optional): Number of votes generated from
            each seed point. Default: 1.
        gt_per_seed (int, optional): Number of ground truth votes generated
            from each seed point. Default: 3.
        num_points (int, optional): Number of points to be used for voting.
            Default: 1.
        conv_channels (tuple[int], optional): Out channels of vote
            generating convolution. Default: (16, 16).
        conv_cfg (dict, optional): Config of convolution.
            Default: dict(type='Conv1d').
        norm_cfg (dict, optional): Config of normalization.
            Default: dict(type='BN1d').
        norm_feats (bool, optional): Whether to normalize features.
            Default: True.
        with_res_feat (bool, optional): Whether to predict residual features.
            Default: True.
        vote_xyz_range (list[float], optional):
            The range of points translation. Default: None.
        vote_loss (dict, optional): Config of vote loss. Default: None.
    """

    def __init__(self, in_channels, vote_per_seed=1, gt_per_seed=3, num_points=-1, conv_channels=(16, 16), conv_cfg=dict(type='Conv1d'), norm_cfg=dict(type='BN1d'), act_cfg=dict(type='ReLU'), norm_feats=True, with_res_feat=True, vote_xyz_range=None, vote_loss=None):
        super().__init__()
        self.in_channels = in_channels
        self.vote_per_seed = vote_per_seed
        self.gt_per_seed = gt_per_seed
        self.num_points = num_points
        self.norm_feats = norm_feats
        self.with_res_feat = with_res_feat
        assert vote_xyz_range is None or is_tuple_of(vote_xyz_range, float)
        self.vote_xyz_range = vote_xyz_range
        if vote_loss is not None:
            self.vote_loss = build_loss(vote_loss)
        prev_channels = in_channels
        vote_conv_list = list()
        for k in range(len(conv_channels)):
            vote_conv_list.append(ConvModule(prev_channels, conv_channels[k], 1, padding=0, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg, bias=True, inplace=True))
            prev_channels = conv_channels[k]
        self.vote_conv = nn.Sequential(*vote_conv_list)
        if with_res_feat:
            out_channel = (3 + in_channels) * self.vote_per_seed
        else:
            out_channel = 3 * self.vote_per_seed
        self.conv_out = nn.Conv1d(prev_channels, out_channel, 1)

    def forward(self, seed_points, seed_feats):
        """forward.

        Args:
            seed_points (torch.Tensor): Coordinate of the seed
                points in shape (B, N, 3).
            seed_feats (torch.Tensor): Features of the seed points in shape
                (B, C, N).

        Returns:
            tuple[torch.Tensor]:

                - vote_points: Voted xyz based on the seed points
                    with shape (B, M, 3), ``M=num_seed*vote_per_seed``.
                - vote_features: Voted features based on the seed points with
                    shape (B, C, M) where ``M=num_seed*vote_per_seed``,
                    ``C=vote_feature_dim``.
        """
        if self.num_points != -1:
            assert self.num_points < seed_points.shape[1], f'Number of vote points ({self.num_points}) should be smaller than seed points size ({seed_points.shape[1]})'
            seed_points = seed_points[:, :self.num_points]
            seed_feats = seed_feats[..., :self.num_points]
        batch_size, feat_channels, num_seed = seed_feats.shape
        num_vote = num_seed * self.vote_per_seed
        x = self.vote_conv(seed_feats)
        votes = self.conv_out(x)
        votes = votes.transpose(2, 1).view(batch_size, num_seed, self.vote_per_seed, -1)
        offset = votes[:, :, :, 0:3]
        if self.vote_xyz_range is not None:
            limited_offset_list = []
            for axis in range(len(self.vote_xyz_range)):
                limited_offset_list.append(offset[..., axis].clamp(min=-self.vote_xyz_range[axis], max=self.vote_xyz_range[axis]))
            limited_offset = torch.stack(limited_offset_list, -1)
            vote_points = (seed_points.unsqueeze(2) + limited_offset).contiguous()
        else:
            vote_points = (seed_points.unsqueeze(2) + offset).contiguous()
        vote_points = vote_points.view(batch_size, num_vote, 3)
        offset = offset.reshape(batch_size, num_vote, 3).transpose(2, 1)
        if self.with_res_feat:
            res_feats = votes[:, :, :, 3:]
            vote_feats = (seed_feats.transpose(2, 1).unsqueeze(2) + res_feats).contiguous()
            vote_feats = vote_feats.view(batch_size, num_vote, feat_channels).transpose(2, 1).contiguous()
            if self.norm_feats:
                features_norm = torch.norm(vote_feats, p=2, dim=1)
                vote_feats = vote_feats.div(features_norm.unsqueeze(1))
        else:
            vote_feats = seed_feats
        return vote_points, vote_feats, offset

    def get_loss(self, seed_points, vote_points, seed_indices, vote_targets_mask, vote_targets):
        """Calculate loss of voting module.

        Args:
            seed_points (torch.Tensor): Coordinate of the seed points.
            vote_points (torch.Tensor): Coordinate of the vote points.
            seed_indices (torch.Tensor): Indices of seed points in raw points.
            vote_targets_mask (torch.Tensor): Mask of valid vote targets.
            vote_targets (torch.Tensor): Targets of votes.

        Returns:
            torch.Tensor: Weighted vote loss.
        """
        batch_size, num_seed = seed_points.shape[:2]
        seed_gt_votes_mask = torch.gather(vote_targets_mask, 1, seed_indices).float()
        seed_indices_expand = seed_indices.unsqueeze(-1).repeat(1, 1, 3 * self.gt_per_seed)
        seed_gt_votes = torch.gather(vote_targets, 1, seed_indices_expand)
        seed_gt_votes += seed_points.repeat(1, 1, self.gt_per_seed)
        weight = seed_gt_votes_mask / (torch.sum(seed_gt_votes_mask) + 1e-06)
        distance = self.vote_loss(vote_points.view(batch_size * num_seed, -1, 3), seed_gt_votes.view(batch_size * num_seed, -1, 3), dst_weight=weight.view(batch_size * num_seed, 1))[1]
        vote_loss = torch.sum(torch.min(distance, dim=1)[0])
        return vote_loss


class Single3DRoIPointExtractor(nn.Module):
    """Point-wise roi-aware Extractor.

    Extract Point-wise roi features.

    Args:
        roi_layer (dict): The config of roi layer.
    """

    def __init__(self, roi_layer=None):
        super(Single3DRoIPointExtractor, self).__init__()
        self.roi_layer = self.build_roi_layers(roi_layer)

    def build_roi_layers(self, layer_cfg):
        """Build roi layers using `layer_cfg`"""
        cfg = layer_cfg.copy()
        layer_type = cfg.pop('type')
        assert hasattr(ops, layer_type)
        layer_cls = getattr(ops, layer_type)
        roi_layers = layer_cls(**cfg)
        return roi_layers

    def forward(self, feats, coordinate, batch_inds, rois):
        """Extract point-wise roi features.

        Args:
            feats (torch.FloatTensor): Point-wise features with
                shape (batch, npoints, channels) for pooling.
            coordinate (torch.FloatTensor): Coordinate of each point.
            batch_inds (torch.LongTensor): Indicate the batch of each point.
            rois (torch.FloatTensor): Roi boxes with batch indices.

        Returns:
            torch.FloatTensor: Pooled features
        """
        rois = rois[..., 1:]
        rois = rois.view(batch_inds, -1, rois.shape[-1])
        with torch.no_grad():
            pooled_roi_feat, pooled_empty_flag = self.roi_layer(coordinate, feats, rois)
            roi_center = rois[:, :, 0:3]
            pooled_roi_feat[:, :, :, 0:3] -= roi_center.unsqueeze(dim=2)
            pooled_roi_feat = pooled_roi_feat.view(-1, pooled_roi_feat.shape[-2], pooled_roi_feat.shape[-1])
            pooled_roi_feat[:, :, 0:3] = rotation_3d_in_axis(pooled_roi_feat[:, :, 0:3], -rois.view(-1, rois.shape[-1])[:, 6], axis=2)
            pooled_roi_feat[pooled_empty_flag.view(-1) > 0] = 0
        return pooled_roi_feat


def get_paddings_indicator(actual_num, max_num, axis=0):
    """Create boolean mask by actually number of a padded tensor.

    Args:
        actual_num (torch.Tensor): Actual number of points in each voxel.
        max_num (int): Max number of points in each voxel

    Returns:
        torch.Tensor: Mask indicates which points are valid inside a voxel.
    """
    actual_num = torch.unsqueeze(actual_num, axis + 1)
    max_num_shape = [1] * len(actual_num.shape)
    max_num_shape[axis + 1] = -1
    max_num = torch.arange(max_num, dtype=torch.int, device=actual_num.device).view(max_num_shape)
    paddings_indicator = actual_num.int() > max_num
    return paddings_indicator


class BaseDGCNNGFModule(nn.Module):
    """Base module for point graph feature module used in DGCNN.

    Args:
        radii (list[float]): List of radius in each knn or ball query.
        sample_nums (list[int]): Number of samples in each knn or ball query.
        mlp_channels (list[list[int]]): Specify of the dgcnn before
            the global pooling for each graph feature module.
        knn_modes (list[str], optional): Type of KNN method, valid mode
            ['F-KNN', 'D-KNN'], Defaults to ['F-KNN'].
        dilated_group (bool, optional): Whether to use dilated ball query.
            Defaults to False.
        use_xyz (bool, optional): Whether to use xyz as point features.
            Defaults to True.
        pool_mode (str, optional): Type of pooling method. Defaults to 'max'.
        normalize_xyz (bool, optional): If ball query, whether to normalize
            local XYZ with radius. Defaults to False.
        grouper_return_grouped_xyz (bool, optional): Whether to return grouped
            xyz in `QueryAndGroup`. Defaults to False.
        grouper_return_grouped_idx (bool, optional): Whether to return grouped
            idx in `QueryAndGroup`. Defaults to False.
    """

    def __init__(self, radii, sample_nums, mlp_channels, knn_modes=['F-KNN'], dilated_group=False, use_xyz=True, pool_mode='max', normalize_xyz=False, grouper_return_grouped_xyz=False, grouper_return_grouped_idx=False):
        super(BaseDGCNNGFModule, self).__init__()
        assert len(sample_nums) == len(mlp_channels), 'Num_samples and mlp_channels should have the same length.'
        assert pool_mode in ['max', 'avg'], "Pool_mode should be one of ['max', 'avg']."
        assert isinstance(knn_modes, list) or isinstance(knn_modes, tuple), 'The type of knn_modes should be list or tuple.'
        if isinstance(mlp_channels, tuple):
            mlp_channels = list(map(list, mlp_channels))
        self.mlp_channels = mlp_channels
        self.pool_mode = pool_mode
        self.groupers = nn.ModuleList()
        self.mlps = nn.ModuleList()
        self.knn_modes = knn_modes
        for i in range(len(sample_nums)):
            sample_num = sample_nums[i]
            if sample_num is not None:
                if self.knn_modes[i] == 'D-KNN':
                    grouper = QueryAndGroup(radii[i], sample_num, use_xyz=use_xyz, normalize_xyz=normalize_xyz, return_grouped_xyz=grouper_return_grouped_xyz, return_grouped_idx=True)
                else:
                    grouper = QueryAndGroup(radii[i], sample_num, use_xyz=use_xyz, normalize_xyz=normalize_xyz, return_grouped_xyz=grouper_return_grouped_xyz, return_grouped_idx=grouper_return_grouped_idx)
            else:
                grouper = GroupAll(use_xyz)
            self.groupers.append(grouper)

    def _pool_features(self, features):
        """Perform feature aggregation using pooling operation.

        Args:
            features (torch.Tensor): (B, C, N, K)
                Features of locally grouped points before pooling.

        Returns:
            torch.Tensor: (B, C, N)
                Pooled features aggregating local information.
        """
        if self.pool_mode == 'max':
            new_features = F.max_pool2d(features, kernel_size=[1, features.size(3)])
        elif self.pool_mode == 'avg':
            new_features = F.avg_pool2d(features, kernel_size=[1, features.size(3)])
        else:
            raise NotImplementedError
        return new_features.squeeze(-1).contiguous()

    def forward(self, points):
        """forward.

        Args:
            points (Tensor): (B, N, C) input points.

        Returns:
            List[Tensor]: (B, N, C1) new points generated from each graph
                feature module.
        """
        new_points_list = [points]
        for i in range(len(self.groupers)):
            new_points = new_points_list[i]
            new_points_trans = new_points.transpose(1, 2).contiguous()
            if self.knn_modes[i] == 'D-KNN':
                idx = self.groupers[i](new_points[..., -3:].contiguous(), new_points[..., -3:].contiguous())[-1]
                grouped_results = grouping_operation(new_points_trans, idx)
                grouped_results -= new_points_trans.unsqueeze(-1)
            else:
                grouped_results = self.groupers[i](new_points, new_points)
            new_points = new_points_trans.unsqueeze(-1).repeat(1, 1, 1, grouped_results.shape[-1])
            new_points = torch.cat([grouped_results, new_points], dim=1)
            new_points = self.mlps[i](new_points)
            new_points = self._pool_features(new_points)
            new_points = new_points.transpose(1, 2).contiguous()
            new_points_list.append(new_points)
        return new_points


class DGCNNGFModule(BaseDGCNNGFModule):
    """Point graph feature module used in DGCNN.

    Args:
        mlp_channels (list[int]): Specify of the dgcnn before
            the global pooling for each graph feature module.
        num_sample (int, optional): Number of samples in each knn or ball
            query. Defaults to None.
        knn_mode (str, optional): Type of KNN method, valid mode
            ['F-KNN', 'D-KNN']. Defaults to 'F-KNN'.
        radius (float, optional): Radius to group with.
            Defaults to None.
        dilated_group (bool, optional): Whether to use dilated ball query.
            Defaults to False.
        norm_cfg (dict, optional): Type of normalization method.
            Defaults to dict(type='BN2d').
        act_cfg (dict, optional): Type of activation method.
            Defaults to dict(type='ReLU').
        use_xyz (bool, optional): Whether to use xyz as point features.
            Defaults to True.
        pool_mode (str, optional): Type of pooling method.
            Defaults to 'max'.
        normalize_xyz (bool, optional): If ball query, whether to normalize
            local XYZ with radius. Defaults to False.
        bias (bool | str, optional): If specified as `auto`, it will be decided
            by the norm_cfg. Bias will be set as True if `norm_cfg` is None,
            otherwise False. Defaults to 'auto'.
    """

    def __init__(self, mlp_channels, num_sample=None, knn_mode='F-KNN', radius=None, dilated_group=False, norm_cfg=dict(type='BN2d'), act_cfg=dict(type='ReLU'), use_xyz=True, pool_mode='max', normalize_xyz=False, bias='auto'):
        super(DGCNNGFModule, self).__init__(mlp_channels=[mlp_channels], sample_nums=[num_sample], knn_modes=[knn_mode], radii=[radius], use_xyz=use_xyz, pool_mode=pool_mode, normalize_xyz=normalize_xyz, dilated_group=dilated_group)
        for i in range(len(self.mlp_channels)):
            mlp_channel = self.mlp_channels[i]
            mlp = nn.Sequential()
            for i in range(len(mlp_channel) - 1):
                mlp.add_module(f'layer{i}', ConvModule(mlp_channel[i], mlp_channel[i + 1], kernel_size=(1, 1), stride=(1, 1), conv_cfg=dict(type='Conv2d'), norm_cfg=norm_cfg, act_cfg=act_cfg, bias=bias))
            self.mlps.append(mlp)


class AllReduce(Function):

    @staticmethod
    def forward(ctx, input):
        input_list = [torch.zeros_like(input) for k in range(dist.get_world_size())]
        dist.all_gather(input_list, input, async_op=False)
        inputs = torch.stack(input_list, dim=0)
        return torch.sum(inputs, dim=0)

    @staticmethod
    def backward(ctx, grad_output):
        dist.all_reduce(grad_output, async_op=False)
        return grad_output


class BasePointSAModule(nn.Module):
    """Base module for point set abstraction module used in PointNets.

    Args:
        num_point (int): Number of points.
        radii (list[float]): List of radius in each ball query.
        sample_nums (list[int]): Number of samples in each ball query.
        mlp_channels (list[list[int]]): Specify of the pointnet before
            the global pooling for each scale.
        fps_mod (list[str], optional): Type of FPS method, valid mod
            ['F-FPS', 'D-FPS', 'FS'], Default: ['D-FPS'].
            F-FPS: using feature distances for FPS.
            D-FPS: using Euclidean distances of points for FPS.
            FS: using F-FPS and D-FPS simultaneously.
        fps_sample_range_list (list[int], optional):
            Range of points to apply FPS. Default: [-1].
        dilated_group (bool, optional): Whether to use dilated ball query.
            Default: False.
        use_xyz (bool, optional): Whether to use xyz.
            Default: True.
        pool_mod (str, optional): Type of pooling method.
            Default: 'max_pool'.
        normalize_xyz (bool, optional): Whether to normalize local XYZ
            with radius. Default: False.
        grouper_return_grouped_xyz (bool, optional): Whether to return
            grouped xyz in `QueryAndGroup`. Defaults to False.
        grouper_return_grouped_idx (bool, optional): Whether to return
            grouped idx in `QueryAndGroup`. Defaults to False.
    """

    def __init__(self, num_point, radii, sample_nums, mlp_channels, fps_mod=['D-FPS'], fps_sample_range_list=[-1], dilated_group=False, use_xyz=True, pool_mod='max', normalize_xyz=False, grouper_return_grouped_xyz=False, grouper_return_grouped_idx=False):
        super(BasePointSAModule, self).__init__()
        assert len(radii) == len(sample_nums) == len(mlp_channels)
        assert pool_mod in ['max', 'avg']
        assert isinstance(fps_mod, list) or isinstance(fps_mod, tuple)
        assert isinstance(fps_sample_range_list, list) or isinstance(fps_sample_range_list, tuple)
        assert len(fps_mod) == len(fps_sample_range_list)
        if isinstance(mlp_channels, tuple):
            mlp_channels = list(map(list, mlp_channels))
        self.mlp_channels = mlp_channels
        if isinstance(num_point, int):
            self.num_point = [num_point]
        elif isinstance(num_point, list) or isinstance(num_point, tuple):
            self.num_point = num_point
        elif num_point is None:
            self.num_point = None
        else:
            raise NotImplementedError('Error type of num_point!')
        self.pool_mod = pool_mod
        self.groupers = nn.ModuleList()
        self.mlps = nn.ModuleList()
        self.fps_mod_list = fps_mod
        self.fps_sample_range_list = fps_sample_range_list
        if self.num_point is not None:
            self.points_sampler = Points_Sampler(self.num_point, self.fps_mod_list, self.fps_sample_range_list)
        else:
            self.points_sampler = None
        for i in range(len(radii)):
            radius = radii[i]
            sample_num = sample_nums[i]
            if num_point is not None:
                if dilated_group and i != 0:
                    min_radius = radii[i - 1]
                else:
                    min_radius = 0
                grouper = QueryAndGroup(radius, sample_num, min_radius=min_radius, use_xyz=use_xyz, normalize_xyz=normalize_xyz, return_grouped_xyz=grouper_return_grouped_xyz, return_grouped_idx=grouper_return_grouped_idx)
            else:
                grouper = GroupAll(use_xyz)
            self.groupers.append(grouper)

    def _sample_points(self, points_xyz, features, indices, target_xyz):
        """Perform point sampling based on inputs.

        If `indices` is specified, directly sample corresponding points.
        Else if `target_xyz` is specified, use is as sampled points.
        Otherwise sample points using `self.points_sampler`.

        Args:
            points_xyz (Tensor): (B, N, 3) xyz coordinates of the features.
            features (Tensor): (B, C, N) features of each point.
            indices (Tensor): (B, num_point) Index of the features.
            target_xyz (Tensor): (B, M, 3) new_xyz coordinates of the outputs.

        Returns:
            Tensor: (B, num_point, 3) sampled xyz coordinates of points.
            Tensor: (B, num_point) sampled points' index.
        """
        xyz_flipped = points_xyz.transpose(1, 2).contiguous()
        if indices is not None:
            assert indices.shape[1] == self.num_point[0]
            new_xyz = gather_points(xyz_flipped, indices).transpose(1, 2).contiguous() if self.num_point is not None else None
        elif target_xyz is not None:
            new_xyz = target_xyz.contiguous()
        elif self.num_point is not None:
            indices = self.points_sampler(points_xyz, features)
            new_xyz = gather_points(xyz_flipped, indices).transpose(1, 2).contiguous()
        else:
            new_xyz = None
        return new_xyz, indices

    def _pool_features(self, features):
        """Perform feature aggregation using pooling operation.

        Args:
            features (torch.Tensor): (B, C, N, K)
                Features of locally grouped points before pooling.

        Returns:
            torch.Tensor: (B, C, N)
                Pooled features aggregating local information.
        """
        if self.pool_mod == 'max':
            new_features = F.max_pool2d(features, kernel_size=[1, features.size(3)])
        elif self.pool_mod == 'avg':
            new_features = F.avg_pool2d(features, kernel_size=[1, features.size(3)])
        else:
            raise NotImplementedError
        return new_features.squeeze(-1).contiguous()

    def forward(self, points_xyz, features=None, indices=None, target_xyz=None):
        """forward.

        Args:
            points_xyz (Tensor): (B, N, 3) xyz coordinates of the features.
            features (Tensor, optional): (B, C, N) features of each point.
                Default: None.
            indices (Tensor, optional): (B, num_point) Index of the features.
                Default: None.
            target_xyz (Tensor, optional): (B, M, 3) new coords of the outputs.
                Default: None.

        Returns:
            Tensor: (B, M, 3) where M is the number of points.
                New features xyz.
            Tensor: (B, M, sum_k(mlps[k][-1])) where M is the number
                of points. New feature descriptors.
            Tensor: (B, M) where M is the number of points.
                Index of the features.
        """
        new_features_list = []
        new_xyz, indices = self._sample_points(points_xyz, features, indices, target_xyz)
        for i in range(len(self.groupers)):
            grouped_results = self.groupers[i](points_xyz, new_xyz, features)
            new_features = self.mlps[i](grouped_results)
            if isinstance(self.mlps[i][0], PAConv):
                assert isinstance(new_features, tuple)
                new_features = new_features[0]
            new_features = self._pool_features(new_features)
            new_features_list.append(new_features)
        return new_xyz, torch.cat(new_features_list, dim=1), indices


class PointSAModuleMSG(BasePointSAModule):
    """Point set abstraction module with multi-scale grouping (MSG) used in
    PointNets.

    Args:
        num_point (int): Number of points.
        radii (list[float]): List of radius in each ball query.
        sample_nums (list[int]): Number of samples in each ball query.
        mlp_channels (list[list[int]]): Specify of the pointnet before
            the global pooling for each scale.
        fps_mod (list[str], optional): Type of FPS method, valid mod
            ['F-FPS', 'D-FPS', 'FS'], Default: ['D-FPS'].
            F-FPS: using feature distances for FPS.
            D-FPS: using Euclidean distances of points for FPS.
            FS: using F-FPS and D-FPS simultaneously.
        fps_sample_range_list (list[int], optional): Range of points to
            apply FPS. Default: [-1].
        dilated_group (bool, optional): Whether to use dilated ball query.
            Default: False.
        norm_cfg (dict, optional): Type of normalization method.
            Default: dict(type='BN2d').
        use_xyz (bool, optional): Whether to use xyz.
            Default: True.
        pool_mod (str, optional): Type of pooling method.
            Default: 'max_pool'.
        normalize_xyz (bool, optional): Whether to normalize local XYZ
            with radius. Default: False.
        bias (bool | str, optional): If specified as `auto`, it will be
            decided by `norm_cfg`. `bias` will be set as True if
            `norm_cfg` is None, otherwise False. Default: 'auto'.
    """

    def __init__(self, num_point, radii, sample_nums, mlp_channels, fps_mod=['D-FPS'], fps_sample_range_list=[-1], dilated_group=False, norm_cfg=dict(type='BN2d'), use_xyz=True, pool_mod='max', normalize_xyz=False, bias='auto'):
        super(PointSAModuleMSG, self).__init__(num_point=num_point, radii=radii, sample_nums=sample_nums, mlp_channels=mlp_channels, fps_mod=fps_mod, fps_sample_range_list=fps_sample_range_list, dilated_group=dilated_group, use_xyz=use_xyz, pool_mod=pool_mod, normalize_xyz=normalize_xyz)
        for i in range(len(self.mlp_channels)):
            mlp_channel = self.mlp_channels[i]
            if use_xyz:
                mlp_channel[0] += 3
            mlp = nn.Sequential()
            for i in range(len(mlp_channel) - 1):
                mlp.add_module(f'layer{i}', ConvModule(mlp_channel[i], mlp_channel[i + 1], kernel_size=(1, 1), stride=(1, 1), conv_cfg=dict(type='Conv2d'), norm_cfg=norm_cfg, bias=bias))
            self.mlps.append(mlp)


class PointSAModule(PointSAModuleMSG):
    """Point set abstraction module with single-scale grouping (SSG) used in
    PointNets.

    Args:
        mlp_channels (list[int]): Specify of the pointnet before
            the global pooling for each scale.
        num_point (int, optional): Number of points.
            Default: None.
        radius (float, optional): Radius to group with.
            Default: None.
        num_sample (int, optional): Number of samples in each ball query.
            Default: None.
        norm_cfg (dict, optional): Type of normalization method.
            Default: dict(type='BN2d').
        use_xyz (bool, optional): Whether to use xyz.
            Default: True.
        pool_mod (str, optional): Type of pooling method.
            Default: 'max_pool'.
        fps_mod (list[str], optional): Type of FPS method, valid mod
            ['F-FPS', 'D-FPS', 'FS'], Default: ['D-FPS'].
        fps_sample_range_list (list[int], optional): Range of points
            to apply FPS. Default: [-1].
        normalize_xyz (bool, optional): Whether to normalize local XYZ
            with radius. Default: False.
    """

    def __init__(self, mlp_channels, num_point=None, radius=None, num_sample=None, norm_cfg=dict(type='BN2d'), use_xyz=True, pool_mod='max', fps_mod=['D-FPS'], fps_sample_range_list=[-1], normalize_xyz=False):
        super(PointSAModule, self).__init__(mlp_channels=[mlp_channels], num_point=num_point, radii=[radius], sample_nums=[num_sample], norm_cfg=norm_cfg, use_xyz=use_xyz, pool_mod=pool_mod, fps_mod=fps_mod, fps_sample_range_list=fps_sample_range_list, normalize_xyz=normalize_xyz)


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)
    (ChamferDistance,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4]), torch.rand([4, 4, 4])], {}),
     False),
    (ConvBNPositionalEncoding,
     lambda: ([], {'input_channel': 4}),
     lambda: ([torch.rand([4, 4, 4])], {}),
     True),
]

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

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

