# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Operator labels and args default value."""
from mindspore.common import dtype as mstype

op_args_default_value = {
    "AdamWeightDecay": {"use_locking": False},
    "AddExt": {"alpha": 1},
    "ApplyCamePart2": {"sum_r": None, "global_shape": None},
    "ApplyCamePart3": {"global_shape": None, "use_first_moment": False},
    "ApplyCamePart4": {"global_shape": None},
    "ApplyRotaryPosEmb": {"cos_format": 0},
    "Argmax": {"axis": -1, "output_type": mstype.int32},
    "ArgMaxWithValue": {"axis": 0, "keep_dims": False},
    "Argmin": {"axis": -1, "output_type": mstype.int32},
    "ArgMinWithValue": {"axis": 0, "keep_dims": False},
    "AvgPoolGrad": {"kernel_size": 1, "strides": 1, "pad_mode": 'VALID', "data_format": 'NCHW'},
    "AvgPool": {"kernel_size": 1, "strides": 1, "pad_mode": 'VALID', "data_format": 'NCHW'},
    "BatchNormGradGrad": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
    "BatchNormGrad": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
    "BatchNormGradWithActivation": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
    "BatchNorm": {"is_training": False, "epsilon": 1e-5, "momentum": 0.1, "data_format": 'NCHW'},
    "BatchNormWithActivation": {"is_training": False, "epsilon": 1e-5, "momentum": 0.1, "data_format": 'NCHW'},
    "BatchNormWithAddAndActivation": {"is_training": False, "epsilon": 1e-5, "momentum": 0.1, "data_format": 'NCHW'},
    "BiasAddGrad": {"data_format": 'NCHW'},
    "BiasAdd": {"data_format": 'NCHW'},
    "CeLU": {"alpha": 1.0},
    "CholeskyInverse": {"upper": False},
    "Cholesky": {"upper": False},
    "Concat": {"axis": 0},
    "Correlate": {"mode": 'valid'},
    "CumProd": {"exclusive": False, "reverse": False},
    "CumSum": {"exclusive": False, "reverse": False},
    "DCT": {"axis": -1, "norm": 'BACKWARD', "forward": True, "grad": False},
    "Diagonal": {"offset": 0, "dim1": 0, "dim2": 1},
    "Dropout": {"keep_prob": 0.5, "Seed0": 0, "Seed1": 0},
    "Eig": {"compute_v": False},
    "Elu": {"alpha": 1.0},
    "ExtractImagePatches": {"padding": 'VALID'},
    "FFT2": {"s": None, "dim": (-2, -1), "norm": None},
    "FFT": {"n": None, "dim": -1, "norm": None},
    "FFTWithSize": {"norm": 'backward', "onesided": True, "signal_sizes": ()},
    "FFTN": {"s": None, "dim": None, "norm": None},
    "FFTShift": {"dim": None},
    "Gather": {"batch_dims": 0},
    "GridSampler2DGrad": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
    "GridSampler2D": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
    "GridSampler3DGrad": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
    "GridSampler3D": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
    "HShrinkGrad": {"lambd": 0.5},
    "HShrink": {"lambd": 0.5},
    "IFFT2": {"s": None, "dim": (-2, -1), "norm": None},
    "IFFT": {"n": None, "dim": -1, "norm": None},
    "IFFTN": {"s": None, "dim": None, "norm": None},
    "IFFTShift": {"dim": None},
    "IRFFTGrad": {"n": None, "dim": -1, "norm": None},
    "IRFFT": {"n": None, "dim": -1, "norm": None},
    "LayerNormGradGrad": {"begin_norm_axis": 1, "begin_params_axis": 1},
    "LayerNormGrad": {"begin_norm_axis": 1, "begin_params_axis": 1},
    "LayerNormGradV3": {"begin_norm_axis": 1, "begin_params_axis": 1},
    "LayerNorm": {"begin_norm_axis": 1, "begin_params_axis": 1, "epsilon": 1e-7},
    "LayerNormV3": {"begin_norm_axis": 1, "begin_params_axis": 1, "epsilon": 1e-7},
    "LogSoftmaxGrad": {"axis": -1},
    "LogSoftmax": {"axis": -1},
    "LogitGrad": {"eps": -1.0},
    "Logit": {"eps": -1.0},
    "MaximumGradGrad": {"grad_x": True, "grad_y": True},
    "MaximumGrad": {"grad_x": True, "grad_y": True},
    "MinimumGrad": {"grad_x": True, "grad_y": True},
    "NanToNum": {"nan": None, "posinf": None, "neginf": None},
    "NLLLossGrad": {"reduction": 'mean', "ignore_index": -100},
    "NLLLoss": {"reduction": 'mean', "ignore_index": -100},
    "OneHot": {"axis": -1},
    "PromptKVCache": {"align_mode": 'LEFT'},
    "Qr": {"full_matrices": False},
    "QuantBatchMatmul": {"offset": None, "bias": None, "transpose_x1": False, "transpose_x2": False, "dtype": mstype.float16},
    "RandpermV2": {"seed": 0, "offset": 0, "dtype": mstype.int64},
    "Range": {"maxlen": 1000000},
    "ReduceAll": {"axis": (), "keep_dims": False},
    "ReduceAny": {"axis": (), "keep_dims": False},
    "ReduceMax": {"axis": (), "keep_dims": False},
    "ReduceMean": {"axis": (), "keep_dims": False},
    "ReduceMin": {"axis": (), "keep_dims": False},
    "ReduceProd": {"axis": (), "keep_dims": False},
    "ReduceStd": {"axis": [], "unbiased": True, "keep_dims": False},
    "ReduceSum": {"axis": (), "keep_dims": False, "skip_mode": False},
    "ResizeBicubicGrad": {"align_corners": False, "half_pixel_centers": False},
    "ResizeBicubic": {"align_corners": False, "half_pixel_centers": False},
    "ResizeBilinearGrad": {"align_corners": False, "half_pixel_centers": False},
    "ResizeBilinearV2": {"align_corners": False, "half_pixel_centers": False},
    "ResizeLinear1DGrad": {"coordinate_transformation_mode": 'align_corners'},
    "ResizeLinear1D": {"coordinate_transformation_mode": 'align_corners'},
    "ResizeNearestNeighborGrad": {"align_corners": False, "half_pixel_centers": False},
    "ResizeNearestNeighbor": {"align_corners": False, "half_pixel_centers": False},
    "ResizeNearestNeighborV2Grad": {"align_corners": False, "half_pixel_centers": False},
    "ResizeNearestNeighborV2": {"align_corners": False, "half_pixel_centers": False},
    "RFFTGrad": {"n": None, "dim": -1, "norm": None},
    "RFFT": {"n": None, "dim": -1, "norm": None},
    "ScalarToTensor": {"dtype": None},
    "SequenceConcat": {"axis": 0},
    "SoftmaxBackward": {"dim": -1},
    "Softmax": {"axis": -1},
    "SolveTriangular": {"trans": 0, "lower": False, "unit_diagonal": False},
    "Split": {"axis": 0, "output_num": 1},
    "StridedSlice": {"begin_mask": 0, "end_mask": 0, "ellipsis_mask": 0, "new_axis_mask": 0, "shrink_axis_mask": 0},
    "SubExt": {"alpha": 1},
    "TupleToTensor": {"dtype": None},
    "WeightQuantBatchMatmul": {"antiquant_offset": None, "quant_scale": None, "quant_offset": None, "bias": None, "transpose_x": False, "transpose_weight": False, "antiquant_group_size": 0},
}

op_labels = {
    "AdamWeightDecay": {"side_effect_mem": True},
    "AssignAdd": {"side_effect_mem": True},
    "Assign": {"side_effect_mem": True},
    "DecoderKVCache": {"side_effect_mem": True},
    "Dropout": {"side_effect_hidden": True},
    "Log": {"cust_aicpu": 'Log', "base": -1.0, "scale": 1.0, "shift": 0.0},
    "PromptKVCache": {"side_effect_mem": True},
    "ReshapeAndCache": {"side_effect_mem": True},
}
