import math
import numbers
import mindspore
import mindtorch
import numpy as np
from mindspore._c_expression import _empty_instance
from ..configs import ENABLE_PYBOOST
from .._op_prim.ascend import legacy, pyboost
from .._op_prim import cpu


def empty(size, dtype):
    return _empty_instance(size, dtype=dtype, device='Ascend')

def empty_like(input, dtype):
    return pyboost.empty_like_op(input, dtype, device='npu')

def new_empty(input, size, dtype, device):
    return pyboost.new_empty_op(input, size, dtype, device)


def reshape(x, shape):
    """
    Reshape the input tensor to the given shape.

    Args:
        x (Tensor): The input tensor.
        shape (tuple): The target shape.

    Returns:
        Tensor: The reshaped tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.reshape_op(x, shape)
    else:
        return legacy.reshape(x, shape)

def contiguous(x):
    """
    Returns a contiguous tensor containing the same data as the input tensor.

    Args:
        x (Tensor): The input tensor.

    Returns:
        Tensor: The contiguous tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.contiguous_op(x)
    else:
        return x

def select_ext_view(input, dim, index):
    """
    Selects a slice from the input tensor along the specified dimension.

    Args:
        input (Tensor): The input tensor.
        dim (int): The dimension along which to select the slice.
        index (int): The index of the slice to select.

    Returns:
        Tensor: The selected slice.
    """
    if ENABLE_PYBOOST:
        # input = clone(input)
        return pyboost.select_ext_view_op(input, dim, index)
    else:
        return legacy.select_view(input, index, dim)

def inplace_copy(self, value):
    """
    Copies the data from the given tensor to the current tensor.

    Args:
        value (Tensor): The tensor from which to copy the data.
    """
    if ENABLE_PYBOOST:
        return pyboost.inplace_copy_op(self, value)
    else:
        legacy.assign(self, value)
        return self

py_slice = slice
def slice(input, dim, start, end, step):
    """
    Slices the input tensor along the specified dimension.

    Args:
        input (Tensor): The input tensor.
        dim (int): The dimension along which to slice.
        start (int): The starting index of the slice.
        end (int): The ending index of the slice.
        step (int): The step size of the slice.

    Returns:
        Tensor: The sliced tensor.
    """
    # if step == 1:
    return pyboost.slice_ext_view_op(input, dim, start, end, step)

def inplace_sub(input, other):
    return inplace_copy(input, legacy.sub(input, other))

    # # ndim = input.ndim
    # # begins = [0] * ndim
    # # ends = [i for i in input.shape]
    # # strides = [1] * ndim
    # # begins[dim] = start
    # # ends[dim] = end
    # # strides[dim] = step
    # # print(input.shape)
    # # print(tuple(begins), tuple(ends), tuple(strides))
    # # print(legacy.strided_slice(input, tuple(begins), tuple(ends), tuple(strides), 0, 0, 0, 0, 0))
    # # return legacy.strided_slice(input, tuple(begins), tuple(ends), tuple(strides), 0, 0, 0, 0, 0)
    # if end < 0:
    #     end = input.shape[dim] + end
    
    # # 2. 计算新视图的大小（size）
    # new_size = list(input.size())
    # # 新维度上的长度计算考虑了步长
    # new_size[dim] = (end - start + step - 1) // step  # 向上取整计算元素个数
    # if new_size[dim] <= 0:
    #     raise RuntimeError(f"Calculated size for dimension {dim} is non-positive after slicing.")

    # # 3. 计算新的步长（stride）和存储偏移量（storage_offset）
    # old_strides = input.stride()
    # new_strides = list(old_strides)
    # # 在目标维度上，新步长 = 原步长 * 步长（step）
    # new_strides[dim] = old_strides[dim] * step
    # # 新的存储偏移量 = 原偏移量 + 起始索引 * 目标维度的原步长
    # new_storage_offset = input.storage_offset() + start * old_strides[dim]

    # # 4. 使用 as_strided 创建新视图
    # # 关键：as_strided 通过直接定义新张量的尺寸、步长和存储偏移量来创建一个视图，而不复制数据。
    # sliced_tensor = as_strided_manual(input, size=tuple(new_size), stride=tuple(new_strides), storage_offset=new_storage_offset)

    # return sliced_tensor

def as_strided_manual(self, size, stride, storage_offset=None):
    if len(size) != len(stride):
        raise RuntimeError("mismatch in length of strides and shape.")
    index = np.arange(0, size[0]*stride[0], stride[0])
    for i in np.arange(1, len(size)):
        tmp = np.arange(0, size[i]*stride[i], stride[i])
        index = np.expand_dims(index, -1)
        index = index + tmp
    if storage_offset is not None:
        index = index + storage_offset

    if index.size == 0:
        input_indices = mindspore.Tensor(Tensor_(index.shape, dtype=mindspore.int32))
    else:
        input_indices = mindspore.tensor(index.astype(np.int32))
    out = gather(reshape(self, (-1,)), input_indices, 0, 0)
    return out


def embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq):
    """
    Applies embedding to the input tensor.

    Args:
        input (Tensor): The input tensor.
        weight (Tensor): The embedding weight tensor.
        padding_idx (int): The index of the padding element.
        max_norm (float): The maximum norm of the embedding vectors.
        norm_type (float): The p-norm to use for normalization.
        scale_grad_by_freq (bool): Whether to scale the gradient by frequency.
        sparse (bool): Whether to use sparse gradients.

    Returns:
        Tensor: The embedded tensor.
    """
    return pyboost.embedding_op(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq)

def add(input, other, alpha=1): # pylint: disable=unused-argument
    """
    Adds two tensors element-wise.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.
        alpha (float): The scaling factor for the other tensor.

    Returns:
        Tensor: The result of the addition.
    """
    if alpha == 1.0:
        return legacy.add(input, other)
    return legacy.add(input, legacy.mul(other, alpha))

def layer_norm(input, normalized_shape, weight, bias, eps=1e-5):
    """
    Applies layer normalization to the input tensor.

    Args:
        input (Tensor): The input tensor.
        normalized_shape (tuple): The shape of the input tensor to be normalized.
        weight (Tensor): The weight tensor.
        bias (Tensor): The bias tensor.
        eps (float): The epsilon value for numerical stability.

    Returns:
        Tensor: The normalized tensor.
    """
    if ENABLE_PYBOOST:
        return cast(pyboost.layer_norm_ext_op(input, normalized_shape, weight, bias, eps)[0], input.dtype)
    if weight is not None:
        begin_axis = input.ndim - weight.ndim
    else:
        begin_axis = -1
    return legacy.layer_norm(input, weight, bias, begin_axis, begin_axis, eps)

def expand_dims(input, axis):
    """
    Adds an extra dimension to the input tensor.

    Args:
        input (Tensor): The input tensor.
        axis (int): The axis along which to add the dimension.

    Returns:
        Tensor: The expanded tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.expand_dims_op(input, axis)
    return legacy.expand_dims(input, axis)

def cast(input, dtype):
    """
    Casts the input tensor to the specified data type.

    Args:
        input (Tensor): The input tensor.
        dtype (str): The target data type.

    Returns:
        Tensor: The casted tensor.
    """
    return legacy.cast(input, dtype)

def sub(input, other, alpha=1):
    """
    Subtracts the other tensor from the input tensor.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The tensor to subtract.
        alpha (float): The scale factor for the other tensor.

    Returns:
        Tensor: The result of the subtraction.
    """
    if ENABLE_PYBOOST:
        return pyboost.sub_ext_op(input, other, alpha)
    return legacy.sub(input, legacy.mul(other, alpha))

def mul(input, other):
    """
    Multiplies the input tensor with the other tensor.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The tensor to multiply.

    Returns:
        Tensor: The result of the multiplication.
    """
    if ENABLE_PYBOOST:
        return pyboost.mul_op(input, other)
    return legacy.mul(input, other)

def inplace_mul(self, other):
    return pyboost.inplace_mul_op(self, other)


def dense(input, weight, bias=None):
    """
    Performs a dense (fully connected) operation.

    Args:
        input (Tensor): The input tensor.
        weight (Tensor): The weight tensor.
        bias (Tensor, optional): The bias tensor. Defaults to None.

    Returns:
        Tensor: The result of the dense operation.
    """
    dtype = input.dtype
    input = cast(input, mindspore.float16)
    weight = cast(weight, mindspore.float16)
    out = cast(pyboost.dense_op(input, weight), dtype)
    if bias is not None:
        out = add(out, bias)
    return out
 

def transpose_view(input, dim0, dim1):
    """
    Transposes the input tensor along the specified dimensions.

    Args:
        input (Tensor): The input tensor.
        dim0 (int): The first dimension to transpose.
        dim1 (int): The second dimension to transpose.

    Returns:
        Tensor: The transposed tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.transpose_ext_view_op(input, dim0, dim1)
    ranks = list(range(input.ndim))
    rank0 = ranks[dim0]
    rank1 = ranks[dim1]
    ranks[dim0] = rank1
    ranks[dim1] = rank0
    return legacy.transpose(input, ranks)

def matmul(input, other):
    """
    Performs a matrix multiplication of the input tensor with another tensor.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.

    Returns:
        Tensor: The result of the matrix multiplication.
    """
    dtype = input.dtype
    input = cast(input, mindspore.float16)
    other = cast(other, mindspore.float16)
    return cast(pyboost.matmul_ext_op(input, other), dtype)

def div(input, other):
    """
    Divides the input tensor by another tensor.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.

    Returns:
        Tensor: The result of the division.
    """
    if ENABLE_PYBOOST:
        return pyboost.div_op(input, other)
    return legacy.div(input, other)

def divmod(input, other, rounding_mode):
    """
    Divides the input tensor by another tensor and returns both the quotient and the remainder.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.
        rounding_mode (str): The rounding mode to use.

    Returns:
        Tuple[Tensor, Tensor]: The quotient and the remainder.
    """
    if rounding_mode == 'floor':
        return legacy.floor_div(input, other)
    elif rounding_mode == 'trunc':
        return legacy.truncate_div(input, other)
    else:
        raise ValueError(f'Invalid rounding mode: {rounding_mode}')

def softmax(input, axis=-1):
    """
    Computes the softmax of the input tensor along the specified axis.

    Args:
        input (Tensor): The input tensor.
        axis (int): The axis along which to compute the softmax.

    Returns:
        Tensor: The softmax of the input tensor.
    """
    if ENABLE_PYBOOST:
        return cast(pyboost.softmax_impl(input, axis), input.dtype)
    return legacy.softmax(input, axis)

def permute(input, axes=None):
    """
    Transposes the dimensions of the input tensor according to the specified axes.

    Args:
        input (Tensor): The input tensor.
        axes (Tuple[int]): The axes to transpose.

    Returns:
        Tensor: The transposed tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.transpose_view_op(input, axes)
    return legacy.transpose(input, axes)

def gelu(input, approximate):
    """
    Computes the Gaussian Error Linear Unit (GELU) activation function.

    Args:
        input (Tensor): The input tensor.

    Returns:
        Tensor: The GELU activation of the input tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.gelu_ext_op(input, approximate)
    return legacy.ge_lu(input)

def tanh(input):
    """
    Computes the hyperbolic tangent of the input tensor.

    Args:
        input (Tensor): The input tensor.

    Returns:
        Tensor: The hyperbolic tangent of the input tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.tanh_op(input)
    return legacy.tanh(input)

def broadcast_to(input, shape):
    """
    Broadcasts the input tensor to the specified shape.

    Args:
        input (Tensor): The input tensor.
        shape (Tuple[int]): The shape to broadcast to.

    Returns:
        Tensor: The broadcasted tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.broadcast_to_view_op(input, shape)
    return legacy.broadcast_to(input, shape)

def split_tensor(tensor, split_size_or_sections, dim):
    """
    Splits a tensor into multiple sub-tensors.

    Args:
        tensor (Tensor): The input tensor.
        split_size_or_sections (Union[int, Tuple[int]]): The size or number of sections to split the tensor into.
        dim (int): The dimension along which to split the tensor.

    Returns:
        List[Tensor]: The list of split sub-tensors.
    """
    if ENABLE_PYBOOST:
        return pyboost.split_tensor_op(tensor, split_size_or_sections, dim)
    return legacy.split(tensor, split_size_or_sections, dim)

def squeeze(input, dim):
    """
    Removes dimensions of size 1 from the shape of the input tensor.

    Args:
        input (Tensor): The input tensor.
        dim (Union[int, Tuple[int]]): The dimensions to squeeze.

    Returns:
        Tensor: The squeezed tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.squeeze_impl(input, dim)
    return legacy.squeeze(input, dim)

def zeros(shape, dtype):
    """
    Returns a tensor filled with zeros.

    Args:
        shape (Union[int, Tuple[int]]): The shape of the tensor.
        dtype (str): The data type of the tensor.

    Returns:
        Tensor: The tensor filled with zeros.
    """
    return legacy.zeros(shape, dtype)

def equal(input, other):
    """
    Returns a tensor with boolean values, indicating element-wise equality.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The tensor to compare with.

    Returns:
        Tensor: The tensor with boolean values.
    """
    if ENABLE_PYBOOST:
        return pyboost.equal_ext_op(input, other)
    return legacy.equal(input, other).all()

def eq(input, other):
    """
    Returns a tensor with boolean values, indicating element-wise equality.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The tensor to compare with.

    Returns:
        Tensor: The tensor with boolean values.
    """
    if ENABLE_PYBOOST:
        return pyboost.equal_op(input, other)
    return legacy.equal(input, other)


def sum(input, dim, keepdim, dtype):
    """
    Returns the sum of elements over a specified dimension.

    Args:
        input (Tensor): The input tensor.
        dim (Union[int, Tuple[int]]): The dimensions to sum over.
        keepdim (bool): Whether to keep the dimensions of size one.

    Returns:
        Tensor: The tensor with summed elements.
    """
    if ENABLE_PYBOOST:
        return pyboost.sum_ext_op(input, dim, keepdim, dtype)
    return legacy.reduce_sum(input.astype(dtype), dim, keepdim)

def dropout(input, p, training=True):
    """
    Returns a tensor with dropout applied element-wise.

    Args:
        input (Tensor): The input tensor.
        p (float): The dropout probability.
        seed (int): The random seed.

    Returns:
        Tensor: The tensor with dropout applied.
    """
    if not training or p==0:
        return input
    return legacy.dropout(input, 1-p, 0, 0)

def clone(input):
    """
    Returns a copy of the input tensor.

    Args:
        input (Tensor): The input tensor.

    Returns:
        Tensor: The copied tensor.
    """
    if ENABLE_PYBOOST:
        return pyboost.clone_op(input)
    return legacy.identity(input)

def inplace_normal(input, mean, std, generator):
    """
    Returns a tensor with normal distribution applied element-wise.

    Args:
        input (Tensor): The input tensor.
        mean (float): The mean of the normal distribution.
        std (float): The standard deviation of the normal distribution.
        seed (int): The random seed.

    Returns:
        Tensor: The tensor with normal distribution applied.
    """
    seed, offset = generator._step(12)
    if ENABLE_PYBOOST:
        return pyboost.inplace_normal_op(input, mean, std, seed, offset)
    return legacy.normal(input, mean, std, 0, 0)

def reduce_all(input, dim, keepdim):
    """
    Returns the sum of all elements in the tensor.

    Args:
        input (Tensor): The input tensor.
        dim (int): The dimension to reduce.
        keepdim (bool): Whether to keep the reduced dimension.

    Returns:
        Tensor: The tensor with the sum of all elements.
    """
    if ENABLE_PYBOOST:
        return pyboost.reduce_all_impl(input, dim, keepdim)
    return legacy.reduce_all(input, dim, keepdim)

def masked_fill(input, mask, value):
    """
    Fills elements of the input tensor with the specified value where the mask is True.

    Args:
        input (Tensor): The input tensor.
        mask (Tensor): The mask tensor.
        value (float): The value to fill.

    Returns:
        Tensor: The tensor with elements filled.
    """
    if ENABLE_PYBOOST:
        return pyboost.masked_fill_op(input, mask, value)
    return legacy.masked_fill(input, mask, value)

def inplace_masked_fill(input, mask, value):
    if isinstance(value, numbers.Number):
        return pyboost.inplace_masked_fill_scalar_op(input, mask, value)
    return pyboost.inplace_masked_fill_tensor_op(input, mask, value)


def isin(input, test_elements, assume_unique=False, invert=False):
    """
    Checks if elements of input tensor are in test_elements.

    Args:
        input (Tensor): The input tensor.
        test_elements (Tensor): The tensor to test against.
        assume_unique (bool): If True, assumes that test_elements contains unique elements.
        invert (bool): If True, inverts the result.

    Returns:
        Tensor: The tensor with boolean values indicating whether elements are in test_elements.
    """
    if ENABLE_PYBOOST:
        return pyboost.isin(input, test_elements, assume_unique, invert)
    return legacy.isin(input, test_elements, assume_unique, invert)

def pad_v3(input, new_pad, mode, value=None, contiguous=True):
    if input.dtype == mindspore.bool_:
        input = cast(input, mindspore.int8)
        out = legacy.pad_v3(input, new_pad, int(value), mode, contiguous)
        return cast(out, mindspore.bool_)
    return legacy.pad_v3(input, new_pad, value, mode, contiguous)

def custom_circular_pad(x, pad):

    ndim = x.ndim
    n_pad_dims = len(pad) // 2
    assert n_pad_dims <= ndim, "填充参数超过了张量的维度"

    # 按从最后维度向前处理填充
    for dim in range(ndim-1, ndim-1-n_pad_dims, -1):
        # 当前维度的左右填充量
        idx = 2 * (ndim - 1 - dim)  # 在pad元组中的起始位置
        left_pad = pad[idx]
        right_pad = pad[idx + 1]
        
        if left_pad == 0 and right_pad == 0:
            continue  # 跳过该维度
            
        size = x.shape[dim]  # 当前维度的原始长度
        new_size = left_pad + size + right_pad
        
        # 生成循环索引: (index - left_pad) mod size
        index = fmod_scalar(add(arange(0, new_size, 1, mindspore.int64), new_size - left_pad), size)
        index = (index + x.shape[dim]) % x.shape[dim]
        x = index_select(x, dim, index)

    return x

def pad(input, pad, mode='constant', value=None):
    if isinstance(pad, tuple):
        pad = tuple(p if isinstance(p, int) else p.item() for p in pad)

    if mode == "constant":
        value = 0 if value is None else value
        return constant_pad_nd(input, pad, value)

    new_pad = ()
    for idx, pad_v in enumerate(pad):
        if not isinstance(pad_v, int):
            pad_v = pad_v.item()
        if pad_v < 0:
            dim = input.ndim - 1 - idx // 2
            input = narrow(input, dim, 0, input.shape[dim] + pad_v)
            pad_v = 0
        new_pad += (pad_v,)
    if py_sum(new_pad) == 0:
        return input
    if mode == 'circular':
        return custom_circular_pad(input, pad)
    elif mode == 'reflect':
        return pad_v3(input, new_pad, mode)
    if value is None:
        value = 0
    if mode == "replicate":
        mode = "edge"
        return pad_v3(input, new_pad, mode)
    # if input.dtype.is_floating_point:
    #     value = float(value)
    # elif input.dtype == mindtorch.bool:
    #     value = bool(value)
    # elif input.dtype in [mindtorch.int32, mindtorch.int64]:
    #     value = int(value)

    return pad_v3(input, new_pad, mode, value)

def log_softmax(input, axis=-1, dtype=None):
    """
    Computes the log softmax of the input tensor along the specified axis.

    Args:
        input (Tensor): The input tensor.
        axis (int): The axis along which to compute the log softmax.
        dtype (dtype): The data type of the output tensor.

    Returns:
        Tensor: The tensor with log softmax values.
    """
    if ENABLE_PYBOOST:
        return pyboost.log_softmax_impl(input, axis)
    return legacy.log_softmax(input, axis)

def not_equal(input, other):
    """
    Computes the element-wise comparison of two tensors for inequality.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.

    Returns:
        Tensor: The tensor with boolean values indicating whether elements are not equal.
    """
    if ENABLE_PYBOOST:
        return pyboost.not_equal_op(input, other)
    return legacy.not_equal(input, other)

def chunk(input, chunks, dim=0):
    """
    Splits a tensor into a specified number of chunks.

    Args:
        input (Tensor): The input tensor.
        chunks (int): The number of chunks to split the tensor into.
        dim (int): The dimension along which to split the tensor.

    Returns:
        Tensor: The tensor split into chunks.
    """
    if ENABLE_PYBOOST:
        return pyboost.chunk_op(input, chunks, dim)
    return legacy.split(input, dim, chunks)

def ones(shape, dtype):
    """
    Returns a tensor filled with ones.

    Args:
        shape (tuple): The shape of the tensor.
        dtype (dtype): The data type of the tensor.

    Returns:
        Tensor: The tensor filled with ones.
    """
    return legacy.ones(shape, dtype)

def greater(input, other):
    """
    Returns a tensor with boolean values indicating whether elements in the input tensor are greater than those in the other tensor.

    Args:
        input (Tensor): The input tensor.
        other (Tensor): The other tensor.

    Returns:
        Tensor: The tensor with boolean values indicating whether elements are greater.
    """
    if ENABLE_PYBOOST:
        return pyboost.greater_op(input, other)
    return legacy.greater(input, other)

def randint(low, high, shape, generator, dtype):
    """
    Returns a tensor filled with random integers from low (inclusive) to high (exclusive).

    Args:
        low (int): The lower bound of the range.
        high (int): The upper bound of the range.
        shape (tuple): The shape of the tensor.
        dtype (dtype): The data type of the tensor.

    Returns:

        Tensor: The tensor filled with random integers.
    """
    seed, offset = generator._step(12)

    if ENABLE_PYBOOST:
        return pyboost.randint_op(low, high, shape, seed, offset, dtype)
    value = legacy.uniform_int(shape,
                                mindspore.tensor(low, dtype=mindspore.int32),
                                mindspore.tensor(high, dtype=mindspore.int32), 0, 0)
    return value
    
def nllloss(input, target, weight, reduction, ingore_index):
    if ENABLE_PYBOOST:
        return pyboost.nllloss_impl(input, target, weight, reduction, ingore_index)[0]
    return legacy.nll_loss(input, target, weight, reduction, ingore_index)[0]

def clamp_scalar(value, min_value, max_value):
    if ENABLE_PYBOOST:
        return pyboost.clamp_scalar_op(value, min_value, max_value)
    if min_value is not None:
        value = legacy.maximum(value, min_value)
    if max_value is not None:
        value = legacy.minimum(value, max_value)
    return value

def cumsum(self, dim, dtype):
    if self.dtype == mindspore.int64:
        self = cast(self, mindspore.int32)
    if ENABLE_PYBOOST:
        return pyboost.cumsum_ext_op(self, dim, dtype)
    if self.shape[dim] == 0:
        return mindtorch.tensor([], dtype=self.dtype, device=self.device)
    return legacy.cum_sum(self, dim, False, False)

def reduce_any(input, axis, keepdims):
    if ENABLE_PYBOOST:
        return pyboost.reduce_any_impl(input, axis, keepdims)
    return legacy.reduce_any(input, axis, keepdims)

def concat(tensors, axis):
    if ENABLE_PYBOOST:
        return pyboost.concat_impl(tensors, axis)
    return legacy.concat(tensors, axis)

def gather_d(input, dim, index):
    idx = meshgrid([arange(0, s, 1, mindspore.int64) for s in index.shape], 'ij')
    
    # 替换目标维度的索引
    new_idx = ()
    for ix, i in enumerate(idx):
        if ix == dim:
            new_idx += (index,)
        else:
            new_idx += (i,)
    
    # 使用高级索引提取数据
    return getitem(input, new_idx)


def gather_nd(input, indices):
    return legacy.gather_nd(input, indices)

def greater_equal(input, other):
    if ENABLE_PYBOOST:
        return pyboost.greater_equal_op(input, other)
    return legacy.greater_equal(input, other)

def less(input, other):
    if ENABLE_PYBOOST:
        return pyboost.less_op(input, other)
    return legacy.less(input, other)

def less_equal(input, other):
    if ENABLE_PYBOOST:
        return pyboost.less_equal_op(input, other)
    return legacy.less_equal(input, other)

def select(condition, input, other):
    return legacy.add(mul(condition, input), mul(bitwise_not(condition), other))

def mean(input, dim, keepdim, dtype):
    if ENABLE_PYBOOST:
        return pyboost.mean_ext_op(input, dim, keepdim, dtype)
    return legacy.reduce_mean(input, dim, keepdim)

def index(input, index):
    if ENABLE_PYBOOST:
        return pyboost.index_op(input, index)
    return legacy.index(input, index)

def scatter(input, dim, index, src):
    # return legacy.tensor_scatter_elements(input, index, cast(src, input.dtype), dim, "none")
    return pyboost.scatter_op(input, dim, index, src)

def tril(input, diagonal=0):
    if ENABLE_PYBOOST:
        return pyboost.tril_ext_op(input, diagonal)
    return legacy.tril(input, diagonal)

def triu(input, diagonal=0):
    if ENABLE_PYBOOST:
        return pyboost.triu_impl(input, diagonal)
    return legacy.triu(input, diagonal)

def inplace_index_put(input, indices, values, accumulate):
    if ENABLE_PYBOOST:
        return pyboost.inplace_index_put_op(input, indices, values, accumulate)
    return legacy.tensor_scatter_elements(input, indices, values, accumulate)

def zeros_like(input, dtype):
    if ENABLE_PYBOOST:
        return pyboost.zeros_like_ext_op(input, dtype)
    return legacy.zeros_like(input)

def ones_like(input, dtype):
    if ENABLE_PYBOOST:
        return pyboost.ones_like_ext_op(input, dtype)
    return legacy.ones_like(input)

def tile(input, multiples):
    return legacy.tile(input, multiples)

def arange(start, end, step, dtype):
    out = legacy.range(start, end, step, 100000)
    # out = mindspore.Tensor(np.arange(start, end, step))
    if dtype is not None:
        out = cast(out, dtype)
    return out

def fill_scalar(input, value, dtype):
    if ENABLE_PYBOOST:
        return pyboost.fill_scalar_op(input, value, dtype)
    return legacy.fill(input, value)

def stop_gradient(input):
    return legacy.stop_gradient(input)

def isinf(input):
    if ENABLE_PYBOOST:
        return pyboost.isinf_op(input)
    return legacy.is_inf(input)

def sort(input, dim, descending, stable):
    out = legacy.sort(input, dim, descending)
    return out[0], cast(out[1], mindspore.int64)

def prod(input, axis, keepdims, dtype):
    if ENABLE_PYBOOST:
        return pyboost.prod_ext_op(input, axis, keepdims, dtype)
    return legacy.reduce_prod(input, axis, keepdims)

def isclose(input, other, rtol, atol, equal_nan):
    return legacy.is_close(input, other, rtol, atol, equal_nan)

def argmax(input, axis, keepdims):
    if ENABLE_PYBOOST:
        return pyboost.argmax_ext_op(input, axis, keepdims)
    return legacy.argmax(input, axis, keepdims)

def argmin(input, axis, keepdims):
    if axis is None:
        axis = -1
    return legacy.arg_min_with_value(input, axis, keepdims)[0]


def bmm(input, other):
    dtype = input.dtype
    out = pyboost.bmm_ext_op(cast(input, mindspore.float16), cast(other, mindspore.float16))
    return cast(out, dtype)

def topk_legacy(input, k, sorted):
    out = legacy.top_k(input, k, sorted)
    return out[0], cast(out[1], mindspore.int64)

def topk(input, k, dim, largest, sorted):
    if not largest:
        input = -input
    if dim is None or dim == input.ndim - 1:
        if not largest:
            res = topk_legacy(input, k, sorted)
            values, indices = -res[0], res[1]
            return values, indices
        return topk_legacy(input, k, sorted)
    input = transpose_view(input, dim, input.ndim - 1)
    output = topk_legacy(input, k, sorted)
    values = transpose_view(output[0], dim, input.ndim - 1)
    indices = transpose_view(output[1], dim, input.ndim - 1)
    if not largest:
        res = (-values, indices)
    else:
        res = (values, indices)
    return res

def logical_not(input):
    if ENABLE_PYBOOST:
        return pyboost.logical_not_op(input)
    return legacy.logical_not(input)

def rand(size, generator, dtype):
    seed, offset = generator._step(12)
    if ENABLE_PYBOOST:
        return pyboost.rand_ext_op(size, seed, offset, dtype)
    return legacy.uniform_real(size, 0, 0)

def inplace_uniform(input, from_, to, generator):
    seed, offset = generator._step(12)
    if ENABLE_PYBOOST:
        return pyboost.uniform_ext_op(input, from_, to, seed, offset)

    if input.dtype.is_floating_point:
        out = legacy.uniform_real(input.shape, 0, 0)
        value = legacy.add(legacy.mul(out, (legacy.sub(to, from_))), from_)
    else:
        value = legacy.uniform_int(input.shape,
                                    mindspore.tensor(from_, dtype=mindspore.int32),
                                    mindspore.tensor(to, dtype=mindspore.int32), 0, 0)
    input.assign_value(legacy.cast(value, input.dtype))

def bitwise_or_tensor(input, other):
    if ENABLE_PYBOOST:
        return pyboost.bitwise_or_tensor_op(input, other)
    return legacy.bitwise_or(input, other)

def bitwise_and_tensor(input, other):
    if ENABLE_PYBOOST:
        return pyboost.bitwise_and_tensor_op(input, other)
    return legacy.bitwise_and(input, other)

def bitwise_or_scalar(input, other):
    if ENABLE_PYBOOST:
        return pyboost.bitwise_or_scalar_op(input, other)
    return legacy.bitwise_or(input, other)

def bitwise_not(input):
    return pyboost.bitwise_not_op(input)

def max(input):
    if ENABLE_PYBOOST:
        return pyboost.max_op(input)
    return legacy.reduce_max(input, (), False)

def stack(tensors, axis=0):
    if ENABLE_PYBOOST:
        return pyboost.stack_ext_impl(tensors, axis)
    return legacy.stack(tensors, axis)

def narrow(input, dim, start, length):
    if ENABLE_PYBOOST:
        return pyboost.narrow_op(input, dim, start, length)
    begin = [0] * input.ndim
    size = [i for i in input.shape]
    begin[dim] = start
    size[dim] = length
    return legacy.slice(input, begin, size)

def std(input, dim, correction, keepdim):
    if dim is None:
        dim = ()
    return legacy.reduce_std(input, dim, bool(correction), keepdim)[0]

def log(input):
    if ENABLE_PYBOOST:
        return pyboost.log_op(input)
    return legacy.log(input)

def gather(input_params, input_indices, axis, batch_dim):
    return legacy.gather(input_params, input_indices, axis, batch_dim)

def non_zero_ext(input):
    if ENABLE_PYBOOST:
        return pyboost.non_zero_ext_op(input)
    return legacy.non_zero(input)

def binary_cross_entropy_with_logits(input, target, weight, posWeight, reduction):
    if ENABLE_PYBOOST:
        return pyboost.binary_cross_entropy_with_logits_impl(input, target, weight, posWeight, reduction)
    return legacy.bce_with_logits_loss(input, target, weight, posWeight, reduction)

def rand_like(input, generator, dtype):
    seed, offset = generator._step(123)
    if ENABLE_PYBOOST:
        return pyboost.rand_like_ext_op(input, seed, offset, dtype)
    return rand(input.shape, dtype)

def floor_div(input, other):
    if ENABLE_PYBOOST:
        return pyboost.floor_div_op(input, other)
    return legacy.floor_div(input, other)

def inplace_fill_scalar(input, value):
    if ENABLE_PYBOOST:
        return pyboost.inplace_fill_scalar_op(input, value)
    input.assign_value(fill_scalar(input.shape, value, input.dtype))
    return input

def linalg_vector_norm(x, ord=2, dim=None, keepdim=False, dtype=None):
    if ENABLE_PYBOOST:
        return pyboost.linalg_vector_norm_op(x, ord, dim, keepdim, dtype)

def non_zero(input):
    if ENABLE_PYBOOST:
        return pyboost.non_zero_op(input)
    return legacy.non_zero(input)

def fmod_scalar(input, other):
    if ENABLE_PYBOOST:
        return pyboost.fmod_scalar_op(input, other)
    return legacy.floor_mod(input, other)

def inplace_zero(input):
    if ENABLE_PYBOOST:
        return pyboost.inplace_zero_op(input)
    input.assign_value(zeros(input.shape, input.dtype))
    return input

def mse_loss(input, target, reduction):
    x = square(input - target)
    average_flag = True
    reduce_flag = True
    if reduction == 'sum':
        average_flag = False
    if reduction == 'none':
        reduce_flag = False

    if reduce_flag and average_flag:
        x = mean(x, tuple(range(x.ndim)), False, None)

    if reduce_flag and not average_flag:
        x = sum(x, tuple(range(x.ndim)), False, None)

    return x

def abs(input):
    if ENABLE_PYBOOST:
        return pyboost.abs_op(input)
    return legacy.abs(input)

def bincount(input, weights=None, minlength=0):
    if ENABLE_PYBOOST:
        return pyboost.bincount_ext_op(input, weights, minlength)
    return legacy.bincount(input, minlength, weights)

def bitwise_and_scalar(input, other):
    if ENABLE_PYBOOST:
        return pyboost.bitwise_and_scalar_op(input, other)
    return legacy.bitwise_and(input, other)

def argmax_with_value(input, axis, keep_dims):
    if ENABLE_PYBOOST:
        return pyboost.argmax_with_value_impl(input, axis, keep_dims)
    return legacy.argmax(input, axis, keep_dims)

def index_select(input, dim, index):
    if ENABLE_PYBOOST:
        return pyboost.index_select_op(input, dim, index)
    return legacy.gather(input, index, dim, 0)

def min(input):
    if ENABLE_PYBOOST:
        return pyboost.min_op(input)
    return legacy.reduce_min(input, (), False)

def minimum(input, other):
    if ENABLE_PYBOOST:
        return pyboost.minimum_op(input, other)
    return legacy.minimum(input, other)

def argmin_with_value(input, axis, keep_dims):
    if ENABLE_PYBOOST:
        return pyboost.argmin_with_value_impl(input, axis, keep_dims)
    return legacy.argmin(input, axis, keep_dims)

def flatten(input, start_dim, end_dim):
    if ENABLE_PYBOOST:
        return pyboost.flatten_ext_op(input, start_dim, end_dim)
    if start_dim < 0:
        start_dim = start_dim + input.ndim
    if end_dim < 0:
        end_dim = end_dim + input.ndim
    input_shape = list(input.shape)
    input_shape[start_dim:end_dim] = [-1]
    return legacy.reshape(input, tuple(input_shape))

def conv2d_padding(input, weight, bias=None, stride=1, padding='valid', dilation=1, groups=1, training=True):
    return conv2d_legacy(input, weight, bias, stride, padding, dilation, groups)

def conv2d(input, weight, bias=None, stride=1, padding='valid', dilation=1, groups=1, training=True):
    return conv2d_legacy(input, weight, bias, stride, padding, dilation, groups)

def conv2d_legacy(input, weight, bias=None, stride=1, padding='valid', dilation=1, groups=1):
    pad_mode = 'pad'
    pad = padding
    if isinstance(padding, (tuple, list)):
        pad = (padding[0], padding[0], padding[1], padding[1])
    elif isinstance(padding, int):
        pad = (padding,) * 4
    if not isinstance(padding, (int, tuple, list)):
        pad_mode = padding
        pad = (0,) * 4
    
    if isinstance(stride, int):
        stride = (stride,) * 4

    out_channels = weight.shape[0]
    kernel_size = weight.shape[2:]

    output = legacy.conv2_d(
        input, weight,
        out_channels,
        kernel_size,
        1,#mode=1,
        pad_mode, #pad_mode=pad_mode,
        pad, #pad=pad,
        tuple(stride), #stride=tuple(stride),
        dilation, #dilation=dilation,
        groups, #group=groups,
        "NCHW", #data_format="NCHW"
    )
    if bias is not None:
        output = legacy.bias_add(output, bias, "NCHW")
    return output

def cos(input):
    if ENABLE_PYBOOST:
        return pyboost.cos_op(input)
    return legacy.cos(input)

def pow_tensor_scalar(input, exponent):
    if ENABLE_PYBOOST:
        return pyboost.pow_tensor_scalar_op(input, exponent)
    return legacy.pow(input, exponent)

def sin(input):
    if ENABLE_PYBOOST:
        return pyboost.sin_op(input)
    return legacy.sin(input)

def batch_norm(input, weight, bias, running_mean=None, running_var=None, training=False, momentum=0.1, epsilon=1e-5):
    if running_mean is None:
        running_mean = ones(input.shape[1], dtype=input.dtype)
    if running_var is None:
        running_var = zeros(input.shape[1], dtype=input.dtype)
    if weight is None:
        weight = ones(input.shape[1], dtype=input.dtype)
    if bias is None:
        bias = zeros(input.shape[1], dtype=input.dtype)

    return legacy.batch_norm(input, weight, bias, running_mean, running_var, training, epsilon, momentum, 'NCHW')

def silu(input):
    if ENABLE_PYBOOST:
        return pyboost.silu_op(input)
    return legacy.silu(input)

def rsqrt(input):
    if ENABLE_PYBOOST:
        return pyboost.rsqrt_op(input)
    return legacy.rsqrt(input)

def sqrt(input):
    if ENABLE_PYBOOST:
        return pyboost.sqrt_op(input)
    return legacy.sqrt(input)

def masked_scatter(input, mask, value):
    return legacy.masked_scatter(input, mask, value)

def inplace_masked_scatter(input, mask, value):
    return pyboost.inplace_masked_scatter_op(input, mask, value)

def neg(input):
    return legacy.neg(input)

def log1p(input):
    return log(add(input, 1))

def pow_scalar_tensor(input, scalar):
    if ENABLE_PYBOOST:
        return pyboost.pow_scalar_tensor_op(input, scalar)
    return legacy.pow(input, scalar)

def adaptive_avg_pool2d(input, output_size):
    return legacy.adaptive_avg_pool2_d(input, output_size)


def exp(input):
    if ENABLE_PYBOOST:
        return pyboost.exp_op(input)
    return legacy.exp(input)

def sigmoid(input):
    if ENABLE_PYBOOST:
        return pyboost.sigmoid_op(input)
    return legacy.sigmoid(input)

def constant_pad_nd(input, pad, value=0.0):
    if ENABLE_PYBOOST:
        return pyboost.constant_pad_nd_op(input, pad, value)

def rfft(input, n=None, dim=-1, norm=None):
    if ENABLE_PYBOOST:
        return pyboost.rfft_op(input, n, dim, norm)
    if input.shape[dim] < n:
        pad_inf = (0, n - input.shape[dim])
        pad_dims = (0, 0) * (input.ndim - (dim + 1)) + pad_inf
        input = constant_pad_nd(input, pad_dims)
    else:
        input = narrow(input, dim, 0, n)
    return legacy.fft_with_size(input, input.ndim, False, True, norm, True, ())

def avg_pool2d(input, kernel_size, stride, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None):
    if ENABLE_PYBOOST:
        return pyboost.avg_pool2d_op(input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)

def conj(input):
    if ENABLE_PYBOOST:
        return pyboost.conj_op(input)
    return legacy.conj(input)

def fill_tensor(size, value, dtype):
    if ENABLE_PYBOOST:
        return pyboost.fill_tensor_op(size, value, dtype)
    return legacy.fill_v2(size, value)

def maximum(input, other):
    if ENABLE_PYBOOST:
        return pyboost.maximum_op(input, other)
    return legacy.maximum(input, other)

def irfft(input, n, dim, norm):
    if ENABLE_PYBOOST:
        return pyboost.irfft_op(input, n, dim, norm)
    return legacy.fft_with_size(input, input.ndim, True, True, norm)

def randn(size, generator, dtype):
    if ENABLE_PYBOOST:
        seed, offset = generator._step(12)
        return pyboost.randn_op(size, seed, offset, dtype)
    return cast(legacy.standard_normal(size, 0, 0), dtype)

def avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True):
    if ENABLE_PYBOOST:
        return pyboost.avg_pool1d_op(input, kernel_size, stride, padding, ceil_mode, count_include_pad)
    return legacy.avg_pool1d(input, kernel_size, stride, padding, ceil_mode, count_include_pad)

def pow(input, exponent):
    if ENABLE_PYBOOST:
        return pyboost.pow_op(input, exponent)
    return legacy.pow(input, exponent)

def roll(input, shifts, axis):
    if ENABLE_PYBOOST:
        return pyboost.roll_impl(input, shifts, axis)
    return legacy.roll(input, shifts, axis)

def conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, training=True):
    return conv1d_legacy(input, weight, bias, stride, padding, dilation, groups)

def conv1d_legacy(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    pad_mode = 'pad'
    pad = padding
    if isinstance(padding, tuple):
        pad = (0, 0, padding[0], padding[0])
    elif isinstance(padding, int):
        pad = (0, 0) + (padding,) * 2
    if not isinstance(padding, (int, tuple)):
        pad_mode = padding
        pad = (0,) * 4

    input = expand_dims(input, 2)
    weight = expand_dims(weight, 2)

    output = legacy.conv2_d(
        input, weight,
        weight.shape[0],
        (1, weight.shape[-1]),
        1,#mode=1,
        pad_mode, #pad_mode=pad_mode,
        pad, #pad=pad,
        (1, stride) if isinstance(stride, int) else (1, *stride), #stride=tuple(stride),
        (1, dilation) if isinstance(dilation, int) else (1, *dilation), #dilation=dilation,
        groups, #group=groups,
        "NCHW", #data_format="NCHW"
    )

    if bias is not None:
        output = legacy.bias_add(output, bias, "NCHW")

    output = squeeze(output, 2)
    return output

def conv1d_padding(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    return conv1d_legacy(input, weight, bias, stride, padding, dilation, groups)

def square(input):
    if ENABLE_PYBOOST:
        return pyboost.square_op(input)
    return legacy.square(input)

def lgamma(input):
    return legacy.lgamma(input)

def reverse_v2(input, dims):
    if isinstance(dims, int):
        dims = (dims,)

    for dim in dims:
        idx = arange(input.size(dim) - 1, -1, -1, None)
        input = index_select(input, dim, idx)
    return input

def unique_consecutive(input, return_inverse, return_counts, dim):
    if ENABLE_PYBOOST:
        return pyboost.unique_consecutive_impl(input, return_inverse, return_counts, dim)
    return legacy.unique_consecutive(input, return_inverse, return_counts, dim)

def split_with_size(input, size, dim=0):
    if ENABLE_PYBOOST:
        return pyboost.split_with_size_op(input, size, dim)
    return legacy.split_with_size(input, size, dim)

def softplus(input, beta=1, threshold=20):
    return legacy.softplus(input)

def remainder_tensor_scalar(input, other):
    out = sub(input, mul(floor_div(input, other), other), 1)
    return out

def baddbmm(input, batch1, batch2, alpha=1, beta=1):
    return add(mul(input, beta), mul(bmm(batch1, batch2), alpha))

def floor(input):
    if ENABLE_PYBOOST:
        return pyboost.floor_op(input)
    return legacy.floor(input)


def _deconv_output_length(pad_mode, filter_size, stride_size, dilation_size, padding):
    """Calculate the width and height of output."""
    length = 0
    filter_size = filter_size + (filter_size - 1) * (dilation_size - 1)
    if pad_mode == 'valid':
        if filter_size - stride_size > 0:
            length = filter_size - stride_size
    elif pad_mode == 'pad':
        length = - padding + filter_size - stride_size

    return length

def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
    pad_mode = 'pad'
    pad = padding
    if isinstance(padding, tuple):
        pad = (padding[0], padding[0], padding[1], padding[1])
    elif isinstance(padding, int):
        pad = (padding,) * 4
    if not isinstance(padding, (int, tuple)):
        pad_mode = padding
        pad = (0,) * 4

    if isinstance(dilation, int):
        dilation = (dilation, dilation)

    in_channel, out_channels = weight.shape[0], weight.shape[1] * groups
    kernel_size = weight.shape[2:]
    n, _, h, w = input.shape
    h_add = _deconv_output_length(pad_mode, kernel_size[0], stride[0], dilation[0], pad[0] + pad[1])
    w_add = _deconv_output_length(pad_mode, kernel_size[1], stride[1], dilation[1], pad[2] + pad[3])

    out = legacy.conv2_d_transpose(
        input, weight,
        (n, out_channels, h * stride[0] + h_add, w * stride[1] + w_add),
        out_channels,
        kernel_size,
        pad_mode,
        pad,
        None,
        1,
        stride,
        dilation,
        groups,
        'NCHW'
    )
    if bias is not None:
        out = legacy.bias_add(out, bias, 'NCHW')
    return out



def relu(input):
    if ENABLE_PYBOOST:
        return pyboost.relu_op(input)
    return legacy.re_lu(input)

def _check_maxpool_padding(padding, nd):
    """Calculate maxpool padding before call primitive"""
    if isinstance(padding, int):
        return (0,) * (3 - nd) + (padding,) * nd
    if isinstance(padding, (tuple, list)):
        if len(padding) == 1:
            return (0,) * (3 - nd) + tuple(padding * nd)
        if len(padding) != nd:
            raise ValueError(f"For {cls_name}, the length of padding must equal to {nd}, but got {len(padding)}.")
        return (0,) * (3 - nd) + tuple(padding)
    return padding

def _cal_dilation(dilation, nd):
    """check the dilation"""
    if isinstance(dilation, int):
        return dilation
    if isinstance(dilation, tuple):
        if len(dilation) == 1:
            return dilation[0]
        if len(dilation) == nd:
            return (3 - nd) * (1,) + dilation
        if nd == 1:
            raise ValueError(f"the length of 'dilation' must be 1, but got {len(dilation)}.")
        raise ValueError(f"the length of 'dilation' must be 1 or {nd}, but got {len(dilation)}.")
    raise ValueError(f"the 'dilation' must be int or tuple, but got {type(dilation)}.")


def max_pool2d(input, kernel_size, stride=1, padding=0, dilation=1, ceil_mode=False, return_indices=False):
    # out, indices = legacy.max_pool_with_argmax_v2(input, kernel_size, stride, padding, dilation, ceil_mode)
    if isinstance(kernel_size, tuple):
        kernel_size = (1,) + kernel_size
    elif isinstance(kernel_size, int):
        kernel_size = (1, kernel_size, kernel_size)
    if isinstance(stride, tuple):
        stride = (1,) + stride
    elif isinstance(stride, int):
        stride = (1, stride, stride)
    padding = _check_maxpool_padding(padding, 2)
    dilation = _cal_dilation(dilation, 2)

    input = expand_dims(input, 2)
    out, indices = legacy.max_pool3_d_with_argmax(input, kernel_size, stride, padding,
                                            dilation, ceil_mode, 'NCDHW', mindspore.int64)
    if return_indices:
        return squeeze(out, 2), squeeze(indices, 2)
    return squeeze(out, 2)

def upsample_bilinear2d(input, size=None, scale_factor=None, align_corners=False):
    return legacy.resize_bilinear_v2(input, size, align_corners, not align_corners)

def group_norm(input, num_groups, weight=None, bias=None, eps=1e-5):
    input_shape = input.shape
    N = input_shape[0]
    C = input_shape[1]
    input_reshaped = reshape(input, (1, N * num_groups, -1 if N!=0 else 1))
    outputs = batch_norm(input_reshaped, None, None, None, None, True, 0., eps)
    out = reshape(outputs, input_shape)
    affine_param_shape = [1] * input.ndim
    affine_param_shape[1] = C
    affine_param_shape = tuple(affine_param_shape)
    if weight is not None and bias is not None:
        out = addcmul(reshape(bias, affine_param_shape), out, reshape(weight, affine_param_shape), value=1)

    elif weight is not None:
        out = mul(out, reshape(weight, affine_param_shape))
    elif bias is not None:
        out = add(out, reshape(bias, affine_param_shape))
    return out

def nllloss_2d(input, target, weight, reduction='mean', ignore_index=-100):
    if ENABLE_PYBOOST:
        return pyboost.nllloss_2d_op(input, target, weight, reduction, ignore_index)[0]
    return legacy.nll_loss(input, target, weight, ignore_index, reduction)[0]

def inplace_relu(input):
    if ENABLE_PYBOOST:
        return pyboost.inplace_relu_op(input)
    return legacy.assign(input, legacy.re_lu(input))

def expm1(input):
    if ENABLE_PYBOOST:
        return pyboost.expm1_op(input)
    return legacy.expm1(input)

def upsample_bicubic2d(input, size=None, scale_factor=None, align_corners=False):
    if ENABLE_PYBOOST:
        return pyboost.upsample_bicubic2d_op(input, size, scale_factor, align_corners)
    return legacy.resize_bicubic(input, size, scale_factor, align_corners)

def acos(input):
    if ENABLE_PYBOOST:
        return pyboost.acos_op(input)
    return legacy.acos(input)

def cdist(x1, x2, p):
    return legacy.cdist(x1, x2, float(p))

def unstack_view(input, dim):
    if ENABLE_PYBOOST:
        return pyboost.unstack_ext_view_op(input, dim)
    return legacy.unstack(input, dim, input.shape[dim])

def l1_loss(input, target, reduction='mean'):
    if ENABLE_PYBOOST:
        return pyboost.l1_loss_ext_op(input, target, reduction)
    return legacy.l1(input, target, reduction)

def diag(input, diagonal):
    if ENABLE_PYBOOST:
        return pyboost.diag_ext_op(input, diagonal)
    return legacy.diag(input, diagonal)

def logsigmoid(input):
    output = sigmoid(input)
    ret = log(output)
    return ret

def one_hot(tensor, num_classes):
    if ENABLE_PYBOOST:
        on_value = mindspore.Tensor(1, dtype=tensor.dtype)
        off_value = mindspore.Tensor(0, dtype=tensor.dtype)
        return pyboost.one_hot_ext_impl(tensor, num_classes, on_value, off_value, -1)
    return legacy.one_hot(tensor, num_classes, on_value, off_value, -1)

def var(input, dim=None, correction=1, keepdim=False):
    if dim is None:
        input_mean = mean(input, (), False, None)
    else:
        input_mean = mean(input, dim=dim, keepdim=True, dtype=None)
    
    # 计算与均值的平方差
    squared_diff = pow(sub(input, input_mean, 1), 2) 
    # 计算方差
    if dim is None:
        variance = mean(squared_diff, (), False, None)
        n = input.numel()  # 总元素个数
    else:
        variance = mean(squared_diff, dim=dim, keepdim=keepdim, dtype=None)
        n = input.size(dim)  # 指定维度的元素个数
    
    # 无偏估计校正
    if correction and n > 1:
        variance = mul(variance, (n / (n - 1)))
    
    return variance


def linspace(start, end, steps, dtype=None):
    start = float(start)
    end = float(end)
    return legacy.lin_space(mindspore.Tensor(start), mindspore.Tensor(end), steps)

def masked_select(input, mask):
    if ENABLE_PYBOOST:
        return pyboost.masked_select_op(input, mask)
    return legacy.masked_select(input, mask)

def glu(input, dim=-1):
    a, b = chunk(input, 2, dim)
    gate = sigmoid(b)
    return mul(a, gate)


def scatter_value(input, dim, index, src, reduce='none'):
    if ENABLE_PYBOOST:
        return pyboost.scatter_value_op(input, dim, index, src, reduce)
    return legacy.scatter(input, dim, index, src, reduce)

def inplace_scatter_value(input, dim, index, src):
    return pyboost.inplace_scatter_value_op(input, dim, index, src)

def inplace_scatter_src(input, dim, index, src):
    return pyboost.inplace_scatter_src_op(input, dim, index, src)


def unique_dim(input, sorted, return_inverse, dim):
    if ENABLE_PYBOOST:
        return pyboost.unique_dim_op(input, sorted, return_inverse, dim)
    return legacy.unique_dim(input, sorted, return_inverse, dim)

def inplace_add(input, other, alpha):
    if isinstance(other, numbers.Number):
        other = mindspore.Tensor(other, dtype=input.dtype)

    if ENABLE_PYBOOST:
        return pyboost.inplace_add_ext_op(input, other, alpha)
    return legacy.inplace_add(input, other)

def logsumexp(input, dim, keepdim):
    input_max = legacy.reduce_max(input, dim, True)
    input_exp = exp(sub(input, input_max))
    input_sumexp = sum(input_exp, dim, keepdim, None)
    input_logsumexp = log(input_sumexp)
    if not keepdim:
        input_max = squeeze(input_max, dim)
    return add(input_logsumexp, input_max)

def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity):
    loss, log_alpha = legacy.ctc_loss_v2(log_probs, targets, input_lengths, target_lengths, blank, 'none', zero_infinity)
    if reduction == 'sum':
        loss = sum(loss, (), False, None)
    if reduction == 'mean':
        # input_type = loss.dtype
        # target_length_t = target_lengths.clip(1., None)
        # loss = loss.astype("float32")
        loss = div(loss, target_lengths)
        loss = mean(loss, (), False, None)
        # loss = loss.astype(input_type)
    return (loss, log_alpha)

def inplace_exponential(self, lambd, generator):
    seed, offset = generator._step(12)
    if ENABLE_PYBOOST:
        return pyboost.inplace_exponential_op(self, lambd, seed, offset)
    return legacy.expo(self, lambd, generator)

def im2col(input, kernel_size, dilation=1, padding=0, stride=1):
    out = legacy.im2_col(input, kernel_size, stride, dilation, padding)
    out_shape = out.shape[:1] + (-1,) + out.shape[-1:]
    out = reshape(out, out_shape)
    return out

def upsample_nearest2d(input, output_size, scale_factors):
    if ENABLE_PYBOOST:
        return pyboost.upsample_nearest2d_op(input, output_size, scale_factors)
    return legacy.upsample_nearest2d(input, scale_factor, align_corners)

def addmm(input, mat1, mat2, alpha=1.0, beta=1.0):
    return add(mul(input, beta), mul(matmul(mat1, mat2), alpha))


def meshgrid(input, lambd):
    if ENABLE_PYBOOST:
        return pyboost.meshgrid_impl(input, lambd)
    return legacy.meshgrid(input, lambd)

def adaptive_avg_pool1d(input, output_size):
    if ENABLE_PYBOOST:
        return pyboost.adaptive_avg_pool1d_op(input, output_size)
    return legacy.adaptive_avg_pool1d(input, output_size)

def conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, training=True):
    pad_mode = 'pad'
    pad = padding
    if isinstance(padding, (tuple, list)):
        pad = (padding[0], padding[0], padding[1], padding[1], padding[2], padding[2])
    elif isinstance(padding, int):
        pad = (padding,) * 6
    if not isinstance(padding, (int, tuple, list)):
        pad_mode = padding
        pad = (0,) * 6

    out_channels = weight.shape[0]
    kernel_size = weight.shape[2:]

    output = legacy.conv3_d(input, weight,
                            out_channels,
                            kernel_size,
                            1,
                            pad_mode,
                            pad,
                            tuple(stride),
                            dilation,
                            groups,
                            "NCDHW")
                            
    if bias is not None:
        output = legacy.bias_add(output, bias, 'NCHW')
    return output


def outer(input, other):
    if ENABLE_PYBOOST:
        return pyboost.outer_op(input, other)
    return legacy.outer(input, other)

def addcmul(input, tensor1, tensor2, value=1.0):
    return legacy.add(mul(mul(tensor1, tensor2), value), input)

def prelu(input, weight):
    if ENABLE_PYBOOST:
        return pyboost.prelu_op(input, weight)
    return legacy.p_re_lu(input, weight)

def reciprocal(input):
    if ENABLE_PYBOOST:
        return pyboost.reciprocal_op(input)
    return legacy.reciprocal(input)

def index_add_ext(input, dim, index, source, alpha):
    if alpha != 1:
        source = mul(alpha, source)
    return legacy.index_add(input, cast(index, mindspore.int32), source, dim, True, True)

def polar(abs, angle):
    if ENABLE_PYBOOST:
        return pyboost.polar_op(abs, angle)
    return legacy.polar(abs, angle)

def upsample_linear1d(input, output_size, scale_factor, align_corners=False):
    coordinate_transformation_mode = "align_corners" if align_corners else "half_pixel"
    return cpu.legacy.resize_linear1_d(input, output_size, coordinate_transformation_mode)


def grid_sampler_2d(input, grid, mode='bilinear', padding_mode='zeros', align_corners=False):
    if ENABLE_PYBOOST:
        return pyboost.grid_sampler_2d_impl(input, grid, mode, padding_mode, align_corners)
    return legacy.grid_sampler_2d(input, grid, mode, padding_mode, align_corners)

def pixel_shuffle(input, upscale_factor):
    if ENABLE_PYBOOST:
        return pyboost.pixel_shuffle_op(input, upscale_factor)
    return legacy.pixel_shuffle(input, upscale_factor)

def view_as_complex(input):
    input = clone(input)
    real_part, imag_part = chunk(input, 2, -1)
    return legacy.complex(squeeze(real_part, -1), squeeze(imag_part, -1))

def view_as_real(input):
    return pyboost.real_view_op(input)

def rms_norm(input, normalized_shape, weight, eps=1e-5):
    if eps is None:
        eps = mindtorch.finfo(input.dtype).eps
    if weight is None:
        weight = ones(normalized_shape, dtype=input.dtype)

    input_dtype = input.dtype
    input = cast(input, mindspore.float32)
    variance = mean(pow(input, 2), -1, True, None)
    input = mul(input, rsqrt(add(variance, eps, 1)))
    return mul(weight, cast(input, input_dtype))

def normal_float_float(mean, std, size, dtype, generator):
    seed, offset = generator._step(12)
    if ENABLE_PYBOOST:
        return pyboost.normal_float_float_op(mean, std, size, seed, offset)

def real(input):
    if ENABLE_PYBOOST:
        return pyboost.real_op(input)
    return legacy.real(input)

def imag(input):
    return legacy.imag(input)

def leaky_relu(input, negative_slope):
    if ENABLE_PYBOOST:
        return pyboost.leaky_relu_ext_op(input, negative_slope)
    return legacy.leaky_relu(input, negative_slope)

def ceil(input):
    if ENABLE_PYBOOST:
        return pyboost.ceil_op(input)
    return legacy.ceil(input)

def erf(input):
    if ENABLE_PYBOOST:
        return pyboost.erf_op(input)
    return legacy.erf(input)

def cross(input, other, dim):
    if ENABLE_PYBOOST:
        return pyboost.cross_impl(input, other, dim)
    return legacy.cross(input, other, dim)

def elu(input, alpha):
    if ENABLE_PYBOOST:
        return pyboost.elu_ext_impl(input, alpha)
    return legacy.elu(input, alpha)

def reduce_max(input, axis, keepdims):
    if ENABLE_PYBOOST:
        return pyboost.reduce_max_impl(input, axis, keepdims)
    return legacy.reduce_max(input, axis, keepdims)

def reduce_min(input, axis, keepdims):
    if ENABLE_PYBOOST:
        return pyboost.reduce_min_impl(input, axis, keepdims)
    return legacy.reduce_min(input, axis, keepdims)

def dynamic_rnn(x, w, b, seq_length, init_h, init_c):
    return legacy.dynamic_rnn(x, w, b, seq_length, init_h, init_c,
                              'LSTM', 'UNIDIRECTIONAL', 1, False, 1.0, -1.0, 0, True, 'tanh', 0.0, True)

def nan_to_num(input, nan=0.0, posinf=None, neginf=None):
    return legacy.nan_to_num(input, nan, posinf, neginf)

def round(input, decimals):
    if ENABLE_PYBOOST:
        return pyboost.round_op(input, decimals)
    return legacy.round(input, decimals)

def fftn(input, s=None, dim=None, norm=None):
    if ENABLE_PYBOOST:
        return pyboost.fftn_op(input, s, dim, norm)

def eye(n, m=None, dtype=None):
    if ENABLE_PYBOOST:
        return pyboost.eye_op(n, m, dtype)
    return legacy.eye(n, m, dtype)

def erfinv(input):
    if ENABLE_PYBOOST:
        return pyboost.erfinv_op(input)
    return legacy.erfinv(input)

def logit(input, eps=1e-5):
    return legacy.logit(input, eps)

def bitwise_xor_tensor(input, other):
    if ENABLE_PYBOOST:
        return pyboost.bitwise_xor_tensor_op(input, other)
    return legacy.bitwise_xor(input, other)

def unique2(input, sorted, return_inverse, return_counts):
    if ENABLE_PYBOOST:
        return pyboost.unique2_op(input, sorted, return_inverse, return_counts)
    return legacy.unique(input, sorted, return_inverse, return_counts)

def sign(input):
    if ENABLE_PYBOOST:
        return pyboost.sign_op(input)
    return legacy.sign(input)

def log2(input):
    if ENABLE_PYBOOST:
        return pyboost.log2_op(input)
    return legacy.log2(input)

def bucketize(input, boundaries, right=False):
    epsilon_ = 0. if right else 1.e-6
    boundaries = [boundary + epsilon_ for boundary in boundaries]
    return legacy.bucketize(input, boundaries)

def inplace_fill_diagonal(input, fill_value, wrap):
    if ENABLE_PYBOOST:
        return pyboost.inplace_fill_diagonal_op(input, fill_value, wrap)
    return legacy.fill_diagonal(input, fill_value, wrap)

def clamp_tensor(input, min, max):
    if ENABLE_PYBOOST:
        return pyboost.clamp_tensor_op(input, min, max)

def hswish(input):
    if ENABLE_PYBOOST:
        return pyboost.hswish_op(input)
    return legacy.h_swish(input)

def logical_and(input, other):
    if ENABLE_PYBOOST:
        return pyboost.logical_and_op(input, other)
    return legacy.logical_and(input, other)

def as_strided(input, size, stride, storage_offset):
    if ENABLE_PYBOOST:
        return pyboost.as_strided_op(input, size, stride, storage_offset)
    return legacy.as_strided(input, size, stride, storage_offset)

def relu6(input):
    if ENABLE_PYBOOST:
        return pyboost.relu6_op(input)
    return legacy.re_lu6(input)

def col2im(input, output_size, kernel_size, dilation=1, padding=0, stride=1):
    if ENABLE_PYBOOST:
        return pyboost.col2im_ext_op(input, output_size, kernel_size, dilation, padding, stride)
    return legacy.col2_im(input, mindspore.Tensor(output_size), kernel_size, dilation, padding, stride)

def flash_attention_score(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode):
    if ENABLE_PYBOOST:
        return pyboost.flash_attention_score_impl(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode)
    return legacy.flash_attention_score(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode)

def prompt_flash_attention(query, key, value, attn_mask, actual_seq_lengths, actual_seq_lengths_kv, pse_shift, deq_scale1, quant_scale1, deq_scale2, quant_scale2, quant_offset2, num_heads, scale_value, pre_tokens, next_tokens, input_layout, num_key_value_heads, sparse_mode, inner_precise):
    return pyboost.prompt_flash_attention_impl(query, key, value, attn_mask, actual_seq_lengths, actual_seq_lengths_kv, pse_shift, deq_scale1, quant_scale1, deq_scale2, quant_scale2, quant_offset2, num_heads, scale_value, pre_tokens, next_tokens, input_layout, num_key_value_heads, sparse_mode, inner_precise)

def incre_flash_attention(query, key, value, attn_mask, actual_seq_lengths, pse_shift, dequant_scale1, quant_scale1, dequant_scale2, quant_scale2, quant_offset2, antiquant_scale, antiquant_offset, block_table, kv_padding_size, num_heads, input_layout, scale_value, num_key_value_heads, block_size, inner_precise):
    return pyboost.incre_flash_attention_impl(query, key, value, attn_mask, actual_seq_lengths, pse_shift, dequant_scale1, quant_scale1, dequant_scale2, quant_scale2, quant_offset2, antiquant_scale, antiquant_offset, block_table, kv_padding_size, num_heads, input_layout, scale_value, num_key_value_heads, block_size, inner_precise)

def randperm(n, generator, dtype):
    seed, offset = generator._step(12)  # pylint: disable=protected-access
    if ENABLE_PYBOOST:
        return pyboost.randperm_ext_op(n, seed, offset, dtype)
    return legacy.randperm(n, seed)

def logical_or(input_x, input_y):
    if ENABLE_PYBOOST:
        return pyboost.logical_or_op(input_x, input_y)
    return legacy.logical_or(input_x, input_y)

def dropout2d(input_x, p):
    return legacy.dropout2_d(input_x, p)

def linalg_qr(input_x, mode):
    if ENABLE_PYBOOST:
        return pyboost.linalg_qr_op(input_x, mode)
    full_matrices = 'mode' == 'complete'
    return legacy.qr(input_x, full_matrices)

def bernoulli(input, generator):
    seed, offset = generator._step(12)
    uniform = rand_like(input, generator, input.dtype)
    result = cast(less(uniform, input), input.dtype)
    return result

def multinomial(input, num_samples, replacement, generator):
    seed, offset = generator._step(12)  # pylint: disable=protected-access
    return pyboost.multinomial_ext_op(input, num_samples, replacement, seed, offset)

def right_shift(input, other):
    if isinstance(other, int):
        other = mindspore.Tensor(other, dtype=input.dtype)
    if ENABLE_PYBOOST:
        return pyboost.right_shift_op(input, other)
    return legacy.right_shift(input, other)

def histc(input, bins=100, min=0, max=0):
    if ENABLE_PYBOOST:
        return pyboost.histc_ext_op(input, bins, float(min), float(max))
    return legacy.histogram(input, bins, float(min), float(max))

def dist_comm_barrier(group):
    return pyboost.dist_comm_barrier_op(group)

def new_empty(input, size, dtype):
    return pyboost.new_empty_op(input, size, dtype, 'Ascend')

def new_ones(input, size, dtype):
    return pyboost.new_ones_op(input, size, dtype)

def kl_div(input, target, reduction, log_target):
    return pyboost.kl_div_op(input, target, reduction, log_target)

def repeat_interleave_int(input, repeats, dim, output_size):
    return pyboost.repeat_interleave_int_op(input, repeats, dim, output_size)

def repeat_interleave_tensor(input, repeats, dim, output_size):
    return pyboost.repeat_interleave_tensor_op(input, repeats, dim, output_size)

def triu_indices(row, col, offset, dtype):
    return legacy.triu_indices(row, col, offset, dtype)

def cumprod(input, dim, dtype):
    out = legacy.cum_prod(input, dim, False, False)
    if dtype is not None:
        out = cast(out, dtype)
    return out

def scatter_nd_update(input, indices, updates):
    return legacy.scatter_nd_update(input, indices, updates, True)

def assign(input, value):
    return inplace_copy(input, value)

def strided_slice(input, begin, end, strides, begin_mask=0, end_mask=0, ellipsis_mask=0, new_axis_mask=0, shrink_axis_mask=0):
    return legacy.strided_slice(input, tuple(begin), tuple(end), tuple(strides), begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask)

def tensor_scatter_update(input, indices, updates):
    return legacy.tensor_scatter_update(input, indices, updates)

def lerp(input, end, weight):
    return legacy.lerp(input, end, weight)

def logaddexp(input, other):
    m = maximum(input, other)
    abs_val = abs(sub(input, other))
    exp_val = exp(neg(abs_val))
    y = add(m, log1p(exp_val))
    return y

def reflection_pad_1d(input, padding):
    return pyboost.reflection_pad_1d_op(input, padding)

def replication_pad_1d(input, padding):
    return pyboost.reflection_pad_1d_op(input, padding)

def hardtanh(input, min_val, max_val):
    return pyboost.hardtanh_op(input, min_val, max_val)

def smooth_l1_loss(input, target, beta=1.0, reduction='none'):
    return pyboost.smooth_l1_loss_impl(input, target, beta, reduction)

def repeat_kv(hidden_states, n_rep: int):
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    # if ON_A2:
    #     return mindtorch.repeat_interleave(hidden_states, dim=1, repeats=n_rep)
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = broadcast_to(expand_dims(hidden_states, 2), (batch, num_key_value_heads, n_rep, slen, head_dim))
    return reshape(hidden_states, batch, num_key_value_heads * n_rep, slen, head_dim)


# @mindspore.jit(capture_mode='ast', backend='ms_backend', dynamic=1, fullgraph=True)
def sdpa(query, key, value, attn_mask=None, dropout_p=0.0,
        is_causal=False, scale=None, enable_gqa=False):
    L, S = query.shape[-2], key.shape[-2]
    scale_factor = 1 / math.sqrt(query.shape[-1]) if scale is None else scale

    if enable_gqa:
        key = contiguous(repeat_kv(key, query.shape[-3] // key.shape[-3]))
        value = contiguous(repeat_kv(value, query.shape[-3] // value.shape[-3]))
    
    attn_bias_shape = (L, S) if attn_mask is None else attn_mask.shape
    attn_bias = zeros(attn_bias_shape, dtype=query.dtype)

    if is_causal:
        assert attn_mask is None
        temp_mask = tril(ones((L, S), dtype=mindtorch.bool), diagonal=0)
        attn_bias = masked_fill(attn_bias, logical_not(temp_mask), mindtorch.finfo(attn_bias.dtype).min)

    if attn_mask is not None:
        if attn_mask.dtype == mindtorch.bool:
            if attn_mask.ndim == 3:
                attn_mask = squeeze(attn_mask, 0)
            else:
                attn_mask = attn_mask
            attn_bias = masked_fill(attn_bias, logical_not(attn_mask), mindtorch.finfo(attn_bias.dtype).min)
        else:
            attn_bias = add(attn_mask, attn_bias)

    attn_weight = mul(matmul(query, transpose_view(key, -2, -1)), scale_factor)
    attn_weight = add(attn_weight, attn_bias)
    attn_weight = softmax(attn_weight, -1)
    attn_weight = dropout(attn_weight, dropout_p)
    return matmul(attn_weight, value)


tensor_1d = mindspore.Tensor([0], dtype=mindtorch.int64)
empty_tensor_1d = mindspore.Tensor(shape=(0,), dtype=mindtorch.int64)
empty_tensor_9d = mindspore.Tensor(shape=(0,)*9, dtype=mindtorch.int64)

def _do_select(self, dim: int, index: int, dim_index: int, self_shape: list):
    """call select view operator"""
    if not self_shape:
        raise TypeError("Invalid index of a 0-dim tensor.")
    dim_size = self_shape[dim]
    if index >= dim_size or index < -dim_size:
        raise IndexError(f"Index {index} is out of bounds for dimension {dim_index} with size {dim_size}")
    index = index + dim_size if index < 0 else index
    return select_ext_view(self, dim, index)

def _wrap_index_to_tuple(index):
    """Wrap index to tuple"""
    if isinstance(index, tuple):
        return index
    if isinstance(index, list):
        if len(index) < 32 and any(isinstance(i, (mindtorch.Tensor, list, tuple, py_slice, type(None), type(...))) for i in index):
            return tuple(index)
    return (index,)

def _do_slice(self, dim: int, index: py_slice, self_shape: list):
    """call slice view operator"""
    def _get_index(index, default):
        if index is None:
            return default
        if mindtorch.is_tensor(index):
            index = int(index)
        return index

    if not self_shape:
        raise TypeError("Invalid index of a 0-dim tensor.")
    step = _get_index(index.step, 1)
    if step <= 0:
        raise ValueError("slice step must be positive")
    start = _get_index(index.start, 0)
    end = _get_index(index.stop, self_shape[dim])
    if start == 0 and end == self_shape[dim] and step == 1:
        return self
    return slice(self, dim, start, end, step)


def _count_indexed_dims(indexes):
    """Count indexed dims"""
    count = 0
    for index in indexes:
        if isinstance(index, mindtorch.Tensor):
            if index.dtype == mindtorch.bool:
                count += index.ndim
            else:
                count += 1
        elif not isinstance(index, (type(None), type(...), bool)):
            count += 1
    return count

def _record_tensor_index(index, remain_indexes, dim):
    """Record indexes remained to be used by aclnnIndex/aclnnIndexPut"""
    if len(remain_indexes) > dim:
        remain_indexes[dim] = index
        return remain_indexes

    while dim > len(remain_indexes):
        # use empty_tensor with dim_num 9 to indicate unused dim
        remain_indexes.append(empty_tensor_9d)

    remain_indexes.append(index)
    return remain_indexes


def _process_dim_in_multi_dim_index(prev_result, orig_tensor, index, dim, indexed_dims, dim_index, remain_indexes,
                                    prev_shape):
    """Process dim in multi dim index"""
    if isinstance(index, bool):
        result = expand_dims(prev_result, dim)
        index_for_bool = tensor_1d if index else empty_tensor_1d
        _record_tensor_index(index_for_bool, remain_indexes, dim)
        prev_shape.insert(dim, 1)
        dim += 1
        return result, dim, remain_indexes, prev_shape
    if isinstance(index, int):
        result = _do_select(prev_result, dim, index, dim_index, prev_shape)
        del prev_shape[dim]
        return result, dim, remain_indexes, prev_shape
    if isinstance(index, py_slice):
        result = _do_slice(prev_result, dim, index, prev_shape)
        # current dim in prev_shape will not be used later, ignore it
        dim += 1
        return result, dim, remain_indexes, prev_shape
    if isinstance(index, type(...)):
        dim += (orig_tensor.ndim - indexed_dims)
        return prev_result, dim, remain_indexes, prev_shape
    if index is None:
        result = expand_dims(prev_result, dim)
        prev_shape.insert(dim, 1)
        dim += 1
        return result, dim, remain_indexes, prev_shape
    if isinstance(index, mindtorch.Tensor):
        result = prev_result
        if index.ndim == 0 and index.dtype in (mindtorch.int, mindtorch.long, mindtorch.short, mindtorch.bool):
            if index.dtype in (mindtorch.int, mindtorch.long, mindtorch.short):
                result = _do_select(prev_result, dim, index.item(), dim_index, prev_shape)
                del prev_shape[dim]
                return result, dim, remain_indexes, prev_shape
            # process index with Tensor bool type
            result = expand_dims(prev_result, dim)
            index_for_bool = tensor_1d if index else empty_tensor_1d
            _record_tensor_index(index_for_bool, remain_indexes, dim)
            prev_shape.insert(dim, 1)
            dim += 1
            return result, dim, remain_indexes, prev_shape
        _record_tensor_index(index, remain_indexes, dim)
        dim += 1
        return result, dim, remain_indexes, prev_shape
    raise IndexError(f"Invalid tensor index type {index}")


def _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims):
    """Process indexes in tuple"""
    self_viewed = self
    self_viewed_shape = list(self.shape)
    dim = 0
    for i, index in enumerate(indexes):
        if isinstance(index, (list, tuple, np.ndarray)):
            index_np = np.array(index) if isinstance(index, (list, tuple)) else index
            if index_np.dtype in (np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64,
                                  np.float16, np.float32, np.float64):
                index = mindspore.tensor(index_np, dtype=mindtorch.int64)
            elif index_np.dtype == np.bool_:
                index = mindspore.tensor(index_np, dtype=mindtorch.int64)
            else:
                raise TypeError(f"Index {index} contain unsupported elements")
        self_viewed, dim, remain_indexes, self_viewed_shape = _process_dim_in_multi_dim_index(
            self_viewed, self, index, dim, indexed_dims, i, remain_indexes, self_viewed_shape)
    return self_viewed, remain_indexes


def getitem(self, index):
    """Handle tensor getitem"""
    if isinstance(index, bool):
        self_viewed = expand_dims(self, 0)
        index_for_bool = tensor_1d if index else empty_tensor_1d
        return pyboost.index_op(self_viewed, [index_for_bool])
    if isinstance(index, int):
        return _do_select(self, 0, index, 0, list(self.shape))
    if isinstance(index, py_slice):
        result = _do_slice(self, 0, index, list(self.shape))
        return result
    if index is None:
        return expand_dims(self, 0)
    if isinstance(index, type(...)):
        return self
    indexes = _wrap_index_to_tuple(index)
    indexed_dims = _count_indexed_dims(indexes)
    if self.ndim < indexed_dims:
        raise IndexError(f"too many indices for tensor with dimension size {self.ndim}")
    remain_indexes = []
    self_viewed, remain_indexes = _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims)
    if not remain_indexes:
        return self_viewed

    return pyboost.index_op(self_viewed, remain_indexes)


def setitem(self, index, value):
    """Handle tensor setitem"""
    if not isinstance(value, mindspore.Tensor):
        if isinstance(value, (bool, int, float)):
            value = mindspore.tensor(value, dtype=self.dtype)
        else:
            raise TypeError(f"Can't assign a {type(value)} to a {self.dtype}.")

    if isinstance(index, bool) and index is False:
        return self
    if isinstance(index, type(...)):
        inplace_copy(self, value)
        return self
    if index is None or (isinstance(index, bool) and index is True):
        self_viewed = expand_dims(self, 0)
        inplace_copy(self_viewed, value)
        return self
    if isinstance(index, int):
        self_viewed = _do_select(self, 0, index, 0, list(self.shape))
        inplace_copy(self_viewed, value)
        return self
    if isinstance(index, py_slice):
        self_viewed = _do_slice(self, 0, index, list(self.shape))
        inplace_copy(self_viewed, value)
        return self
    indexes = _wrap_index_to_tuple(index)
    indexed_dims = _count_indexed_dims(indexes)
    if self.ndim < indexed_dims:
        raise IndexError(f"too many indices for tensor with dimension size {self.ndim}")
    remain_indexes = []
    self_viewed, remain_indexes = _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims)
    if not remain_indexes:
        inplace_copy(self_viewed, value)
        return self

    inplace_index_put(self_viewed, remain_indexes, value, False) # accumulate=False
    return self

def isin(elements, test_elements, invert=False):
    elements_shape = elements.shape
    elements = expand_dims(reshape(elements, (-1,)), -1)
    if not isinstance(test_elements, numbers.Number):
        test_elements = reshape(test_elements, (-1,))
    included = eq(elements, test_elements)
    # ops.reduce_sum only supports float
    res = cast(sum(included, -1, False, None), mindtorch.bool_)
    if invert:
        res = logical_not(res)
    return reshape(res, elements_shape)

def mish(input):
    if ENABLE_PYBOOST:
        return pyboost.mish_ext_op(input)
    return legacy.mish(input)

def _get_unfold_indices(input_shape, dimension, size, step):
    if dimension < 0:
        dimension += len(input_shape)
    indices = []
    for i in range(0, input_shape[dimension] - size + 1, step):
        indices.append(list(range(i, i + size)))

    return indices, dimension


def unfold(input, dimension, size, step):
    _indices, _dimension = _get_unfold_indices(input.shape, dimension, size, step)
    indices = mindspore.tensor(_indices)
    output = gather(input, indices, _dimension, 0)
    output = transpose_view(output, _dimension + 1, -1)
    return output

def selu(input):
    """SELU activation: scale * elu(x, alpha) where alpha=1.67326324, scale=1.05070098"""
    SELU_ALPHA = 1.67326324
    SELU_SCALE = 1.05070098
    if ENABLE_PYBOOST:
        return pyboost.mul_op(legacy.elu(input, SELU_ALPHA), SELU_SCALE)
    return legacy.mul(legacy.elu(input, SELU_ALPHA), SELU_SCALE)

def celu(input, alpha):
    """CELU activation: max(0, x) + min(0, alpha * (exp(x/alpha) - 1))"""
    if alpha == 0:
        raise ZeroDivisionError("ZeroDivisionError: alpha cannot be 0 for CELU")
    return elu(input, alpha)

def hardsigmoid(input):
    """Hardsigmoid activation: clamp((x + 3) / 6, 0, 1)"""
    x_plus_3 = add(input, 3.0)
    x_div_6 = div(x_plus_3, 6.0)
    return clamp_scalar(x_div_6, 0.0, 1.0)

def fast_gelu(x):
    """Fast GELU approximation"""
    return gelu(x, approximate='tanh')

def swiglu(x, dim=-1):
    """Swish-Gated Linear Unit: swish(x[..., :d]) * x[..., d:] where d = x.shape[dim] // 2"""
    split_size = x.shape[dim] // 2
    x1, x2 = legacy.split(x, split_size, dim)
    if ENABLE_PYBOOST:
        return pyboost.mul_op(silu(x1), x2)
    return legacy.mul(silu(x1), x2)

def rotary_position_embedding(x, cos, sin, mode=0):
    """Rotary Position Embedding"""
    import mindspore
    return mindspore.ops.auto_generate.gen_ops_def.apply_rotary_pos_emb_(x, cos, sin, mode)
