# The operators included in this file are:
# 1) Operators defined only in PIR, dynamic graphs do not exist;
# 2) The definitions of static graphs and dynamic graphs are inconsistent, but the final definition plan has not yet been clarified.
# After the definition is clearly defined, migrate to paddle /fluid/pir/dialect/operator/ir/update_ops.yaml or paddle/phi/api/yaml/ops.yaml

- op : adadelta_
  args : (Tensor param, Tensor grad, Tensor avg_squared_grad, Tensor avg_squared_update, Tensor learning_rate, Tensor master_param, float rho, float epsilon, bool multi_precision)
  output : Tensor(param_out), Tensor(moment_out), Tensor(inf_norm_out), Tensor(master_param_out)
  infer_meta :
    func : AdadeltaInferMeta
  kernel :
    func : adadelta
    data_type : param
  optional : master_param, master_param_out
  inplace : (param -> param_out), (avg_squared_grad -> moment_out), (avg_squared_update -> inf_norm_out), (master_param -> master_param_out)

- op : add
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : add
  inplace : (x -> out)
  backward : add_grad
  interfaces : paddle::dialect::InferSymbolicShapeInterface

- op : add_n
  args : (Tensor[] inputs)
  output : Tensor
  invoke : add_n_impl(inputs)
  backward : add_n_grad

- op : add_n_
  args : (Tensor[] inputs)
  output : Tensor(out)
  infer_meta:
    func: AddNInferMeta
    param: [inputs]
  kernel:
    func: add_n
    param: [inputs]

- op : add_n_with_kernel
  args : (Tensor[] inputs)
  output : Tensor(out)
  infer_meta:
    func: AddNInferMeta
    param: [inputs]
  kernel:
    func: add_n
    param: [inputs]

- op : all
  args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
  kernel :
    func : all

- op : amax
  args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
  kernel :
    func : amax
  backward : amax_grad

- op : amin
  args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
  kernel :
    func : amin
  backward : amin_grad

- op : any
  args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
  kernel :
    func : any

- op : assign
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : assign
  backward : assign_grad
  inplace : (x -> out)

- op : assign_out_
  args : (Tensor x, Tensor output)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : assign
    param : [x]
  inplace : (output -> out)
  backward : assign_out__grad

- op : assign_value
  args : (int[] shape, DataType dtype, Scalar[] values, Place place = {})
  output : Tensor(out)
  infer_meta :
    func : AssignValueInferMeta
    param: [shape, dtype]
  kernel :
    func : assign_value
    param : [shape, dtype, values]
    backend: place>
    data_type : dtype

- op : assign_value_
  args : (Tensor output, int[] shape, DataType dtype, Scalar[] values, Place place = {})
  output : Tensor(out)
  inplace: (output -> out)
  infer_meta :
    func : AssignValueInferMeta
    param : [shape, dtype]
  kernel :
    func : assign_value
    param : [shape, dtype, values]
    data_type : dtype
    backend : place > output

- op : batch_norm
  args : (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_format, bool use_global_stats, bool trainable_statistics)
  output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)
  infer_meta:
    func : BatchNormInferMeta
  kernel :
    func : batch_norm
    data_type : x
  view : (mean -> mean_out), (variance -> variance_out)
  backward : batch_norm_grad
  optional : scale, bias, reserve_space

- op : c_allgather
  args : (Tensor x, int ring_id, int nranks, bool use_calc_stream)
  output : Tensor(out)
  infer_meta :
    func : AllGatherInferMeta
    param: [x, nranks]
  kernel :
    func : c_allgather

- op : c_allreduce_max
  args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
  output : Tensor(out)
  infer_meta :
    func : AllReduceInferMeta
    param : [x]
  kernel :
    func : c_allreduce_max
  inplace : (x -> out)

- op : c_allreduce_sum
  args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
  output : Tensor(out)
  infer_meta :
    func : AllReduceInferMeta
    param : [x]
  kernel :
    func : c_allreduce_sum
  inplace : (x -> out)

- op : c_broadcast
  args : (Tensor x, int ring_id=0, int root=0, bool use_calc_stream=false)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : c_broadcast
  inplace : (x -> out)

- op : c_concat
  args : (Tensor x, int rank, int nranks, int ring_id, bool use_calc_stream, bool use_model_parallel)
  output : Tensor(out)
  infer_meta :
    func : CConcatInferMeta
    param : [x, nranks]
  kernel :
    func : c_concat

- op : c_embedding
  args : (Tensor weight, Tensor x, int64_t start_index=0, int64_t vocab_size=-1)
  output : Tensor(out)
  infer_meta :
    func : CEmbeddingInferMeta
    param : [weight, x, start_index]
  kernel :
    func : c_embedding
    param : [weight, x, start_index, vocab_size]
    data_type : weight
  backward : c_embedding_grad

- op : c_identity
  args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
  output : Tensor(out)
  infer_meta :
    func : CIdentityInferMeta
  kernel :
    func : c_identity
  inplace : (x -> out)

- op : c_reduce_min
  args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
  output : Tensor(out)
  infer_meta :
    func : DistReduceInferMeta
    param : [x]
  kernel :
    func : c_reduce_min
  inplace : (x -> out)

- op : c_reduce_sum
  args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
  output : Tensor(out)
  infer_meta :
    func : DistReduceInferMeta
    param : [x]
  kernel :
    func : c_reduce_sum
  inplace : (x -> out)

- op : c_reducescatter
  args : (Tensor x, int ring_id = 0, int nranks = 1, bool use_calc_stream = false)
  output : Tensor(out)
  infer_meta :
    func : ReduceScatterInferMeta
    param: [x, nranks]
  kernel :
    func : reduce_scatter
    param: [x, nranks]

- op : c_sync_calc_stream
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : c_sync_calc_stream
  inplace : (x -> out)

- op : c_sync_comm_stream
  args : (Tensor x, int ring_id)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : c_sync_comm_stream
  inplace : (x -> out)

- op : cast
  args : (Tensor x, DataType dtype)
  output : Tensor(out)
  infer_meta :
    func : CastInferMeta
    spmd_rule : CastInferSpmd
  kernel :
    func : cast
    param : [x, dtype]
    data_type : x
  inplace: (x -> out)
  backward : cast_grad
  interfaces : paddle::dialect::InferSymbolicShapeInterface

- op : channel_shuffle
  args : (Tensor x, int groups, str data_format="NCHW")
  output : Tensor(out)
  infer_meta :
    func : ChannelShuffleInferMeta
  kernel :
    func : channel_shuffle
  backward : channel_shuffle_grad

- op : conv2d_transpose
  args : (Tensor x, Tensor filter, int[] strides={1, 1}, int[] paddings={0, 0}, int[] output_padding={}, IntArray output_size={}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1}, str data_format="NCHW")
  output : Tensor(out)
  infer_meta :
    func : Conv2dTransposeInferMeta
  kernel :
    func : conv2d_transpose
    data_type : x
  backward : conv2d_transpose_grad

- op : copy_to
  args : (Tensor x, Place place, bool blocking)
  output : Tensor(out)
  invoke : copy_to_impl(x, place, blocking)

- op : decayed_adagrad
  args : (Tensor param, Tensor grad, Tensor moment, Tensor learning_rate, float decay = 0.95f, float epsilon = 1.0e-6f)
  output : Tensor(param_out), Tensor(moment_out)
  infer_meta :
    func : DecayedAdagradInferMeta
  kernel :
    func : decayed_adagrad
    data_type : param

- op : decode_jpeg
  args : (Tensor x, str mode, Place place)
  output : Tensor(out)
  infer_meta :
    func : DecodeJpegInferMeta
    param : [x, mode]
  kernel :
    func : decode_jpeg
    param : [x, mode]
    backend : place

- op : deformable_conv
  args : (Tensor x, Tensor offset, Tensor filter, Tensor mask, int[] strides, int[] paddings, int[] dilations, int deformable_groups, int groups, int im2col_step)
  output : Tensor(out)
  infer_meta :
    func : DeformableConvInferMeta
  kernel :
    func : deformable_conv
    data_type : x
  optional : mask
  backward : deformable_conv_grad

- op : depthwise_conv2d_transpose
  args : (Tensor x, Tensor filter, int[] strides={1, 1}, int[] paddings={0, 0}, int[] output_padding={}, IntArray output_size={}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1}, str data_format="NCHW")
  output : Tensor(out)
  infer_meta :
    func : Conv2dTransposeInferMeta
  kernel :
    func : depthwise_conv2d_transpose
    data_type : x
  backward : depthwise_conv2d_transpose_grad

- op : disable_check_model_nan_inf
  args: (Tensor x, int flag = 0)
  output: Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param : [x]
  kernel:
    func: check_model_nan_inf
    data_type: x
  backward: disable_check_model_nan_inf_grad

- op : distribute_fpn_proposals
  args : (Tensor fpn_rois, Tensor rois_num, int min_level, int max_level, int refer_level, int refer_scale, bool pixel_offset)
  output : Tensor[](multi_fpn_rois){max_level - min_level + 1}, Tensor[](multi_level_rois_num){max_level - min_level + 1}, Tensor(restore_index)
  infer_meta :
    func : DistributeFpnProposalsInferMeta
  kernel :
    func : distribute_fpn_proposals
    data_type : fpn_rois
  optional : rois_num, multi_level_rois_num

- op : divide
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : divide
  inplace: (x -> out)
  backward : divide_grad

- op : dropout
  args : (Tensor x, Tensor seed_tensor, Scalar p, bool is_test, str mode, int seed, bool fix_seed)
  output : Tensor(out), Tensor(mask)
  infer_meta :
    func : DropoutInferMeta
  kernel :
    func : dropout
    data_type : x
  optional : seed_tensor
  intermediate : mask
  backward : dropout_grad

- op : einsum
  args : (Tensor[] x, str equation)
  output : Tensor(out), Tensor[](inner_cache){x.size()}, Tensor[](xshape){x.size()}
  infer_meta :
    func : EinsumRawInferMeta
    param : [x, equation]
  kernel :
    func : einsum
  optional : inner_cache, xshape
  backward : einsum_grad

- op : elementwise_pow
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule: ElementwiseBinaryInferSpmd
  kernel :
    func : elementwise_pow
  backward : elementwise_pow_grad

- op : embedding
  args : (Tensor x, Tensor weight, int64_t padding_idx=-1, bool sparse=false)
  output : Tensor
  infer_meta :
    func : EmbeddingInferMeta
    param : [x, weight, padding_idx]
  kernel :
    func : embedding {dense, dense -> dense}
           sparse_weight_embedding {dense, selected_rows -> dense}
    param : [x, weight, padding_idx]
    data_type : weight
  backward : embedding_grad

- op : empty
  args : (IntArray shape, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
  output: Tensor(out)
  infer_meta :
    func : CreateInferMeta
    param : [shape, dtype]
  kernel :
    func : empty
    param : [shape, dtype]
    data_type : dtype
    backend : place

- op : empty_like
  args : (Tensor x, DataType dtype = DataType::UNDEFINED, Place place = {})
  output: Tensor(out)
  infer_meta :
    func : CreateLikeInferMeta
    param : [x, dtype]
  kernel :
    func : empty_like
    param : [x, dtype]
    data_type : dtype > x
    backend : place > x

- op : enable_check_model_nan_inf
  args: (Tensor x, int flag = 1)
  output: Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param : [x]
  kernel:
    func: check_model_nan_inf
    data_type: x
  backward: enable_check_model_nan_inf_grad

- op : equal
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
    spmd_rule: ElementwiseBinaryInferSpmd
  kernel :
    func : equal
  inplace: (x -> out)

- op : exponential_
  args : (Tensor x, float lam)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : exponential
  inplace : (x -> out)
  backward : exponential__grad

- op : eye
  args : (Scalar num_rows, Scalar num_columns, DataType dtype=DataType::FLOAT32, Place place={})
  output : Tensor(out)
  infer_meta :
    func : EyeInferMeta
    param : [num_rows, num_columns, dtype]
  kernel :
    func : eye
    param : [num_rows, num_columns, dtype]
    data_type : dtype
    backend : place

- op : feed
  args : (str name, int col)
  output : Tensor(out)

- op : fetch
  args : (Tensor x, str name, int col)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : fetch
    param : [x]
  traits : pir::SideEffectTrait

- op : floor_divide
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
  kernel :
    func : floor_divide
  inplace: (x -> out)

- op : frobenius_norm
  args : (Tensor x, IntArray axis,  bool keep_dim,  bool reduce_all)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMetaBase
  kernel :
    func : frobenius_norm
  backward : frobenius_norm_grad

- op : full
  args : (IntArray shape, Scalar value, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
  output: Tensor(out)
  infer_meta :
    func : CreateInferMeta
    param : [shape, dtype]
  kernel :
    func : full
    param : [shape, value, dtype]
    data_type : dtype
    backend : place

- op : full_
  args : (Tensor output, IntArray shape, Scalar value, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
  output : Tensor(out)
  inplace : (output -> out)
  infer_meta :
    func : CreateInferMeta
    param : [shape, dtype]
  kernel :
    func : full
    param : [shape, value, dtype]
    data_type : dtype
    backend : place

- op : full_batch_size_like
  args : (Tensor input, int[] shape, DataType dtype, Scalar value, int input_dim_idx, int output_dim_idx, Place place=CPUPlace())
  output: Tensor(out)
  infer_meta :
    func : FullBatchSizeLikeInferMeta
    param : [input, shape, value, dtype, input_dim_idx, output_dim_idx]
  kernel :
    func : full_batch_size_like
    param : [input, shape, value, dtype, input_dim_idx, output_dim_idx]
    data_type : dtype
    backend : place

- op : full_like
  args : (Tensor x, Scalar value, DataType dtype = DataType::UNDEFINED, Place place = {})
  output: Tensor(out)
  infer_meta :
    func : CreateLikeInferMeta
    param : [x, dtype]
  kernel :
    func : full_like
    param : [x, value, dtype]
    data_type : dtype > x
    backend : place > x
  data_transform :
    skip_transform : x

- op : full_with_tensor
  args : (Tensor shape, Tensor value, DataType dtype=DataType::FLOAT32)
  output: Tensor(out)
  infer_meta :
    func : FullWithTensorInferMeta
    param : [shape, dtype]
  kernel :
    func : full_with_tensor
    data_type : dtype

- op : fused_adam_
  args : (Tensor[] params, Tensor[] grads, Tensor learning_rate, Tensor[] moments1, Tensor[] moments2, Tensor[] beta1_pows, Tensor[] beta2_pows, Tensor[] master_params, Tensor skip_update, Scalar beta1, Scalar beta2, Scalar epsilon, int chunk_size, float weight_decay, bool use_adamw, bool multi_precision, bool use_global_beta_pow)
  output : Tensor[](params_out){params.size()}, Tensor[](moments1_out){params.size()}, Tensor[](moments2_out){params.size()}, Tensor[](beta1_pows_out){params.size()}, Tensor[](beta2_pows_out){params.size()}, Tensor[](master_params_out){params.size()}
  infer_meta :
    func : FusedAdamInferMeta
  kernel :
    func : fused_adam
    data_type : params
  optional : skip_update, master_params
  inplace : (params -> params_out), (moments1 -> moments1_out), (moments2 -> moments2_out), (beta1_pows -> beta1_pows_out), (beta2_pows -> beta2_pows_out), (master_params -> master_params_out)

- op : fused_batch_norm_act
  args : (Tensor x, Tensor scale, Tensor bias, Tensor mean, Tensor variance, float momentum, float epsilon, str act_type)
  output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)
  infer_meta:
    func : FusedBatchNormActInferMeta
    param : [x, scale, bias, mean, variance]
  kernel :
    func : fused_batch_norm_act
    data_type : x
  view : (mean -> mean_out), (variance -> variance_out)
  backward : fused_batch_norm_act_grad

- op : fused_bn_add_activation
  args : (Tensor x, Tensor z, Tensor scale, Tensor bias, Tensor mean, Tensor variance, float momentum, float epsilon, str act_type)
  output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)
  infer_meta:
    func : FusedBatchNormActInferMeta
    param : [x, scale, bias, mean, variance]
  kernel :
    func : fused_bn_add_activation
    data_type : x
  view : (mean -> mean_out), (variance -> variance_out)
  backward : fused_bn_add_activation_grad

- op : fused_softmax_mask_upper_triangle
  args : (Tensor X)
  output : Tensor(Out)
  infer_meta :
    func : UnchangedInferMeta
  kernel:
    func : fused_softmax_mask_upper_triangle
  backward: fused_softmax_mask_upper_triangle_grad

- op : gaussian
  args : (IntArray shape, float mean, float std, int seed, DataType dtype, Place place={})
  output: Tensor(out)
  infer_meta :
    func : GaussianInferMeta
    param : [shape, mean, std, seed, dtype]
  kernel :
    func : gaussian
    param : [shape, mean, std, seed, dtype]
    data_type : dtype
    backend : place

- op : get_tensor_from_selected_rows
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel:
    func: get_tensor_from_selected_rows {selected_rows -> dense}

- op : greater_equal
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
  kernel :
    func : greater_equal
  inplace: (x -> out)

- op : greater_than
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
  kernel :
    func : greater_than
  inplace: (x -> out)

- op : hardswish
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : hardswish
  backward : hardswish_grad

- op : hsigmoid_loss
  args : (Tensor x, Tensor label, Tensor w, Tensor bias, Tensor path, Tensor code, int num_classes, bool is_sparse)
  output : Tensor(out), Tensor(pre_out), Tensor(w_out)
  infer_meta :
    func : HSigmoidLossInferMeta
  optional: path, code, bias
  kernel :
    func : hsigmoid_loss
    data_type : x
  backward : hsigmoid_loss_grad

- op : increment
  args : (Tensor x, float value = 1.0)
  output : Tensor(out)
  infer_meta :
    func : IncrementInferMeta
  kernel :
    func : increment
  inplace : (x -> out)

- op : less_equal
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
  kernel :
    func : less_equal
  inplace: (x -> out)

- op : less_than
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
  kernel :
    func : less_than
  inplace: (x -> out)

- op : linspace
  args : (Tensor start, Tensor stop, Tensor number, DataType dtype, Place place)
  output : Tensor(out)
  infer_meta :
    func : LinspaceInferMeta
    param: [start, stop, number, dtype]
  kernel :
    func : linspace
    param: [start, stop, number, dtype]
    data_type : dtype
    backend : place

- op : load_combine
  args : (str file_path, bool load_as_fp16, bool model_from_memory)
  output : Tensor[](Out)
  kernel:
    func: load_combine
    param: [file_path, load_as_fp16, model_from_memory]
  optional : Out

- op : lod_array_length
  args : (Tensor[] x)
  output : Tensor(out)

- op : logspace
  args : (Tensor start, Tensor stop, Tensor num, Tensor base, DataType dtype, Place place={})
  output : Tensor(out)
  infer_meta:
    func : LogspaceInferMeta
    param : [start, stop, num, base, dtype]
  kernel :
    func : logspace
    param : [start, stop, num, base, dtype]
    data_type : dtype
    backend : place

- op : logsumexp
  args : (Tensor x, int64_t[] axis,  bool keepdim,  bool reduce_all)
  output : Tensor(out)
  infer_meta :
    func : LogsumexpInferMeta
  kernel :
    func : logsumexp
  backward : logsumexp_grad

- op : lrn
  args : (Tensor x, int n=5, float k=2.0, float alpha=0.0001, float beta=0.75, str data_format="NCHW")
  output : Tensor(out), Tensor(mid_out)
  infer_meta :
    func : LrnInferMeta
    param : [x, n]
  kernel :
    func : lrn
  backward : lrn_grad

- op : matmul
  args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false)
  output : Tensor
  infer_meta :
    func : MatmulInferMeta
    spmd_rule : MatmulInferSpmd
  kernel :
    func : matmul
  backward : matmul_grad

- op : matrix_rank
  args : (Tensor x, float tol, bool use_default_tol=true, bool hermitian=false)
  output : Tensor(out)
  infer_meta :
    func : MatrixRankInferMeta
    param : [x, use_default_tol, hermitian]
  kernel :
    func : matrix_rank

- op : matrix_rank_tol
  args : (Tensor x, Tensor atol_tensor, bool use_default_tol=true, bool hermitian=false)
  output : Tensor(out)
  infer_meta :
    func : MatrixRankTolInferMeta
  kernel :
    func : matrix_rank_tol

- op : max
  args : (Tensor x, IntArray axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMeta
    spmd_rule: ReductionMaxInferSpmdDynamic
  kernel :
    func : max
  backward : max_grad

- op : maximum
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : maximum
  backward : maximum_grad

- op : mean
  args : (Tensor x, IntArray axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMeta
    spmd_rule : ReductionMeanInferSpmdDynamic
  kernel :
    func : mean
  backward : mean_grad

- op : memcpy
  args : (Tensor x, int dst_place_type)
  output : Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param: [x]
  kernel:
    func : memcpy
    param: [x, dst_place_type]

- op : memcpy_d2h
  args : (Tensor x, int dst_place_type)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : memcpy_d2h

- op : memcpy_h2d
  args : (Tensor x, int dst_place_type)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : memcpy_h2d

- op : min
  args : (Tensor x, IntArray axis={}, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMeta
  kernel :
    func : min
  backward : min_grad

- op : minimum
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
  kernel :
    func : minimum
  backward : minimum_grad

- op : mish
  args : (Tensor x, float lambda)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : mish
  backward : mish_grad

- op : multiply
  args : (Tensor x, Tensor y)
  output : Tensor
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : multiply {dense, dense -> dense},
           multiply_sr {selected_rows, dense -> selected_rows}
  inplace : (x -> out)
  backward : multiply_grad

- op : norm
  args : (Tensor x, int axis, float epsilon, bool is_test)
  output : Tensor(out), Tensor(norm)
  infer_meta :
    func : NormInferMeta
  kernel :
    func : norm
  backward : norm_grad

- op : not_equal
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : CompareInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : not_equal
  inplace: (x -> out)

- op : one_hot
  args : (Tensor x, Scalar(int) num_classes)
  output : Tensor(out)
  infer_meta :
    func : OneHotInferMeta
  kernel :
    func : one_hot

- op : ones
  args : (IntArray shape, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
  output : Tensor(out)
  invoke : full(shape, 1, dtype, place)

- op : ones_like
  args : (Tensor x, DataType dtype=DataType::UNDEFINED, Place place={})
  output : Tensor(out)
  invoke : full_like(x, 1, dtype, place)

- op : pad
  args : (Tensor x, int[] paddings, Scalar pad_value)
  output : Tensor
  infer_meta :
    func : PadInferMeta
  kernel :
    func : pad
  backward : pad_grad

- op : pool2d
  args : (Tensor x, IntArray kernel_size, int[] strides, int[] paddings, bool ceil_mode, bool exclusive, str data_format, str pooling_type, bool global_pooling, bool adaptive, str padding_algorithm)
  output : Tensor(out)
  infer_meta :
    func : Pool2DInferMeta
    param : [x, kernel_size, strides, paddings, ceil_mode, exclusive, data_format, pooling_type, global_pooling, adaptive, padding_algorithm]
  kernel :
    func : pool2d
    param : [x, kernel_size, strides, paddings, ceil_mode, exclusive, data_format, pooling_type, global_pooling, adaptive, padding_algorithm]
  backward : pool2d_grad

- op : pool3d
  args : (Tensor x, int[] kernel_size, int[] strides, int[] paddings, bool ceil_mode, bool exclusive, str data_format, str pooling_type, bool global_pooling, bool adaptive, str padding_algorithm)
  output : Tensor(out)
  infer_meta :
    func : PoolInferMeta
    param : [x, kernel_size, strides, paddings, ceil_mode, exclusive, data_format, pooling_type, global_pooling, adaptive, padding_algorithm]
  kernel :
    func : pool3d
    param : [x, kernel_size, strides, paddings, ceil_mode, exclusive, data_format, pooling_type, global_pooling, adaptive, padding_algorithm]
  backward : pool3d_grad

- op : print
  args : (Tensor in, int first_n, str message, int summarize, bool print_tensor_name = true, bool print_tensor_type = true, bool print_tensor_shape = true, bool print_tensor_layout = true, bool print_tensor_lod = true, str print_phase = "BOTH", bool is_forward = true)
  output : Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param: [in]
  kernel :
    func : print_kernel
    param: [in, first_n, message, summarize, print_tensor_name, print_tensor_type, print_tensor_shape, print_tensor_layout, print_tensor_lod, print_phase, is_forward]

- op : prod
  args : (Tensor x, IntArray dims, bool keep_dim, bool reduce_all)
  output : Tensor
  infer_meta :
    func : ReduceIntArrayAxisInferMetaBase
  kernel :
    func : prod
  backward : prod_grad

- op : randint
  args : (int low, int high, IntArray shape, DataType dtype=DataType::INT64, Place place={})
  output : Tensor(out)
  infer_meta :
    func : RandintInferMeta
    param : [low, high, shape, dtype]
  kernel :
    func : randint
    param : [low, high, shape, dtype]
    data_type : dtype
    backend : place

- op : randperm
  args : (int n, DataType dtype, Place place={})
  output : Tensor(out)
  infer_meta :
    func : RandpermInferMeta
    param : [n, dtype]
  kernel :
    func : randperm
    param : [n, dtype]
    data_type : dtype
    backend : place

- op : read_file
  args : (str filename = "", DataType dtype=DataType::UINT8, Place place=CPUPlace())
  output : Tensor(out)
  infer_meta :
    func : ReadFileInferMeta
    param : [filename]
  kernel :
    func : read_file
    param : [filename]
    data_type : dtype
    backend : place

- op : recv_v2
  args : (int[] out_shape = {}, DataType dtype = DataType::FLOAT32, int peer = 0, int ring_id = 0, bool use_calc_stream = false, bool dynamic_shape = false)
  output : Tensor(out)
  infer_meta:
    func: RecvV2InferMeta
    param: [ring_id, dynamic_shape, peer, out_shape, dtype]
  kernel :
    func : recv_v2
    param : [ring_id, dynamic_shape, peer, out_shape, dtype, use_calc_stream]
    data_type : dtype

- op : remainder
  args : (Tensor x, Tensor y)
  output : Tensor (out)
  infer_meta :
    func : ElementwiseInferMeta
  kernel :
    func : remainder
  inplace : (x -> out)

- op : repeat_interleave
  args : (Tensor x, int repeats, int axis)
  output : Tensor(out)
  infer_meta :
    func : RepeatInterleaveInferMeta
  kernel :
    func : repeat_interleave
    data_type : x
  backward: repeat_interleave_grad

- op : repeat_interleave_with_tensor_index
  args : (Tensor x, Tensor repeats, int axis)
  output : Tensor(out)
  infer_meta :
    func : RepeatInterleaveWithTensorIndexInferMeta
  kernel :
    func : repeat_interleave_with_tensor_index
    data_type : x
  backward: repeat_interleave_with_tensor_index_grad

- op : reshape
  args : (Tensor x, IntArray shape)
  output : Tensor(out), Tensor(xshape)
  infer_meta :
    func : ReshapeWithXShapeInferMeta
  kernel :
    func : reshape
  inplace : (x -> out)
  view: (x -> out)
  intermediate : xshape
  backward: reshape_grad
  interfaces : paddle::dialect::InferSymbolicShapeInterface

- op : rnn
  args: (Tensor x, Tensor[] pre_state, Tensor[] weight_list, Tensor sequence_length, Tensor dropout_state_in, float dropout_prob=0.0, bool is_bidirec=false, int input_size=10, int hidden_size=100, int num_layers=1, str mode="RNN_TANH", int seed=0, bool is_test=false)
  output: Tensor(out), Tensor(dropout_state_out), Tensor[](state){pre_state.size()}, Tensor(reserve)
  infer_meta:
    func: RnnInferMeta
    param : [x, pre_state, weight_list, sequence_length, dropout_prob, is_bidirec, input_size, hidden_size, num_layers, mode, seed, is_test]
  kernel:
    func: rnn
    param : [x, pre_state, weight_list, sequence_length, dropout_prob, is_bidirec, input_size, hidden_size, num_layers, mode, seed, is_test]
    data_type: x
  backward: rnn_grad
  optional : sequence_length
  intermediate : reserve
  view : (dropout_state_in -> dropout_state_out)

- op : row_conv
  args : (Tensor x, Tensor filter)
  output : Tensor(out)
  infer_meta :
     func: RowConvInferMeta
  kernel :
     func : row_conv

- op : rrelu
  args : (Tensor x, float lower, float upper, bool is_test)
  output : Tensor(out), Tensor(noise)
  infer_meta :
    func : RReluInferMeta
  kernel :
    func : rrelu
    data_type : x
  intermediate : noise
  backward : rrelu_grad

- op : save_combine
  args : (Tensor[] x, str file_path, bool overwrite, bool save_as_fp16, bool save_to_memory)
  output : Tensor(out)
  kernel:
    func: save_combine_tensor
    param: [x, file_path, overwrite, save_as_fp16, save_to_memory]
  optional : out

- op : seed
  args : (int seed, bool deterministic, str rng_name, bool force_cpu)
  output : Tensor(out)
  infer_meta:
    func: SeedInferMeta
    param: [seed]
  kernel:
    func: seed

- op : send_v2
  args : (Tensor x, int ring_id = 0, int peer = 0, bool use_calc_stream = false, bool dynamic_shape = false)
  output :
  infer_meta:
    func: SendV2InferMeta
    param: [peer, ring_id]
  kernel :
    func : send_v2
    param : [x, ring_id, dynamic_shape, peer, use_calc_stream]

- op : set_value
  args : (Tensor x, IntArray starts, IntArray ends, IntArray steps, int64_t[] axes, int64_t[] decrease_axes, int64_t[] none_axes, int64_t[] shape, Scalar[] values)
  output : Tensor(out)
  inplace: (x -> out)
  infer_meta :
    func : SetValueInferMeta
    param : [x]
  kernel :
    func : set_value
  backward: set_value_grad

- op : set_value_with_tensor
  args : (Tensor x, Tensor values, IntArray starts, IntArray ends, IntArray steps, int64_t[] axes, int64_t[] decrease_axes, int64_t[] none_axes)
  output : Tensor(out)
  inplace: (x -> out)
  infer_meta:
    func: SetValueInferMeta
    param: [x]
  kernel:
    func: set_value_with_tensor
  backward: set_value_with_tensor_grad

- op : shadow_feed
  args : (Tensor x)
  output : Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param: [x]
  kernel:
    func: shadow_feed
    param: [x]

- op : share_data
  args : (Tensor x)
  output : Tensor(out)
  infer_meta:
    func: UnchangedInferMeta
    param: [x]
  kernel:
    func: share_data
    param: [x]

- op : shuffle_batch
  args : (Tensor x, Tensor seed, int startup_seed=0)
  output : Tensor(out), Tensor(shuffle_idx), Tensor(seed_out)
  infer_meta:
     func: ShuffleBatchInferMeta
  kernel:
     func: shuffle_batch
     data_type: x
  backward : shuffle_batch_grad

- op : slice
  args : (Tensor input, int64_t[] axes, IntArray starts, IntArray ends, int64_t[] infer_flags, int64_t[] decrease_axis)
  output : Tensor
  infer_meta :
    func : SliceRawInferMeta
    spmd_rule : SliceInferSpmdDynamic
  kernel :
    func : slice
  backward : slice_grad
  interfaces : paddle::dialect::InferSymbolicShapeInterface

- op : soft_relu
  args : (Tensor x, float threshold = 20.0f)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : soft_relu
  backward : soft_relu_grad

- op : softmax
  args : (Tensor x, int axis)
  output : Tensor(out)
  infer_meta :
    func : SoftmaxInferMeta
  kernel :
    func : softmax
  inplace : (x -> out)
  backward : softmax_grad

- op : split
  args : (Tensor x, IntArray sections, Scalar(int) axis)
  output : Tensor[]{sections.size()}
  infer_meta :
    func : SplitInferMeta
  kernel :
    func : split
  backward : split_grad

- op : split_with_num
  args : (Tensor x, int num, Scalar(int) axis)
  output : Tensor[]{num}
  infer_meta :
    func : SplitWithNumInferMeta
    spmd_rule : SplitWithNumInferSpmdDynamic
  kernel :
    func : split_with_num
  backward : split_with_num_grad

- op : strided_slice
  args : (Tensor x, int[] axes, IntArray starts, IntArray ends, IntArray strides)
  output : Tensor
  infer_meta :
    func : StridedSliceInferMeta
    spmd_rule : StridedSliceInferSpmdDynamic
  kernel :
    func : strided_slice
  backward : strided_slice_grad

- op : subtract
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseInferMeta
    spmd_rule : ElementwiseBinaryInferSpmd
  kernel :
    func : subtract
  inplace : (x -> out)
  backward : subtract_grad
  interfaces : paddle::dialect::InferSymbolicShapeInterface

- op : sum
  args : (Tensor x, IntArray axis={}, DataType dtype=DataType::UNDEFINED, bool keepdim=false)
  output : Tensor(out)
  infer_meta :
    func : SumInferMeta
    spmd_rule : ReductionSumInferSpmdDynamic
  kernel :
    func : sum
    data_type : x
  backward : sum_grad

- op : swish
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : swish
  backward : swish_grad

- op : sync_batch_norm_
  args : (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_format, bool use_global_stats, bool trainable_statistics)
  output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)
  infer_meta :
    func : BatchNormInferMeta
  kernel :
    func : sync_batch_norm
    data_type : x
  backward : sync_batch_norm_grad
  inplace : (mean -> mean_out), (variance -> variance_out)
  optional : reserve_space

- op : tile
  args : (Tensor x, IntArray repeat_times = {})
  output : Tensor(out)
  infer_meta :
    func : TileInferMeta
  kernel :
    func : tile
  backward : tile_grad

- op : trans_layout
  args : (Tensor x, int[] perm)
  output : Tensor
  infer_meta :
    func : TransposeInferMeta
  kernel :
    func : transpose
  backward : trans_layout_grad

- op : transpose
  args : (Tensor x, int[] perm)
  output : Tensor(out)
  infer_meta :
    func : TransposeInferMeta
    spmd_rule: TransposeInferSpmd
  kernel :
    func : transpose
  inplace : (x -> out)
  backward : transpose_grad

- op : tril
  args : (Tensor x, int diagonal)
  output : Tensor(out)
  infer_meta :
    func : TrilInferMeta
  kernel :
    func : tril
  inplace: (x -> out)
  backward : tril_grad

- op : tril_indices
  args : (int rows, int cols, int offset, DataType dtype, Place place={})
  output : Tensor(out)
  infer_meta :
    func : TrilIndicesInferMeta
    param : [rows, cols, offset, dtype]
  kernel :
    func : tril_indices
    param : [rows, cols, offset, dtype]
    data_type : dtype
    backend : place

- op : triu
  args : (Tensor x, int diagonal)
  output : Tensor(out)
  infer_meta :
    func : TriuInferMeta
    spmd_rule : TriuInferSpmd
  kernel :
    func : triu
  inplace: (x -> out)
  backward : triu_grad

- op : triu_indices
  args : (int row, int col, int offset, DataType dtype, Place place={})
  output : Tensor(out)
  infer_meta :
    func : TriuIndicesInferMeta
    param : [row, col, offset, dtype]
  kernel :
    func : triu_indices
    param : [row, col, offset, dtype]
    data_type : dtype
    backend : place

# python API: paddle.nn.initializer.TruncatedNormal
- op : truncated_gaussian_random
  args : (int[] shape, float mean, float std, int seed, DataType dtype=DataType::FLOAT32, Place place={})
  output : Tensor(out)
  infer_meta :
    func : TruncatedGaussianRandomInferMeta
    param : [shape, mean, std, seed, dtype]
  kernel :
    func : truncated_gaussian_random
    param : [shape, mean, std, seed, dtype]
    backend : place
    data_type : dtype

- op : uniform
  args : (IntArray shape,  DataType dtype,  Scalar min,  Scalar max,  int seed, Place place={})
  output : Tensor(out)
  infer_meta :
    func : UniformRandomInferMeta
    param: [shape, dtype]
  kernel :
    func : uniform
    param: [shape, dtype, min, max, seed]
    data_type : dtype
    backend : place

- op : uniform_random_batch_size_like
  args : (Tensor input, int[] shape, int input_dim_idx=0, int output_dim_idx=0, float min=-1.0f, float max=1.0f, int seed=0, int diag_num=0,  int diag_step=0, float diag_val=1.0f, DataType dtype=DataType::FLOAT32)
  output : Tensor(out)
  infer_meta :
    func : BatchSizeLikeInferMeta
    param : [input,shape,input_dim_idx,output_dim_idx]
  kernel :
    func : uniform_random_batch_size_like
    data_type : dtype

- op : unique
  args : (Tensor x, bool return_index=false, bool return_inverse=false, bool return_counts=false, int[] axis={}, DataType dtype=DataType::INT64, bool is_sorted=false)
  output : Tensor(out), Tensor(indices), Tensor(inverse), Tensor(counts)
  optional : indices, counts
  infer_meta :
    func : UniqueRawInferMeta
  kernel :
    func : unique
    data_type : x
  interfaces : paddle::dialect::ParseKernelKeyInterface

- op : unpool
  args: (Tensor x, Tensor indices, int[] ksize, int[] strides, int[] padding, IntArray output_size, str data_format)
  output: Tensor(out)
  infer_meta:
    func: UnpoolInferMeta
  kernel:
    func: unpool
    data_type: x
  backward: unpool_grad

- op : write_to_array
  args : (Tensor i, Tensor x)
  output : Tensor[](out)

- op : zeros
  args : (IntArray shape, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
  output : Tensor(out)
  invoke : full(shape, 0, dtype, place)

- op : zeros_like
  args : (Tensor x, DataType dtype=DataType::UNDEFINED, Place place = {})
  output : Tensor(out)
  invoke : full_like(x, 0, dtype, place)

- op: c_softmax_with_cross_entropy
  args: (Tensor logits, Tensor label,  int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0)
  output: Tensor(softmax), Tensor(loss)
  infer_meta:
    func : CSoftmaxWithCrossEntropyInferMeta
  kernel:
    func: c_softmax_with_cross_entropy
    data_type : logits
  backward: c_softmax_with_cross_entropy_grad

- op: dpsgd
  args: (Tensor param, Tensor grad, Tensor learning_rate, float clip = 10.0f, float batch_size = 16.0f, float sigma = 1.0f, int seed = 0)
  output: Tensor(param_out)
  infer_meta:
     func: DpsgdInferMeta
  kernel:
     func: dpsgd
     data_type: param

- op: ftrl
  args: (Tensor param, Tensor squared_accumulator, Tensor linear_accumulator, Tensor grad, Tensor learning_rate, float l1=0.0f, float l2=0.0f, float lr_power=-0.5f)
  output: Tensor(param_out), Tensor(squared_accum_out), Tensor(linear_accum_out)
  infer_meta:
    func: FtrlInferMeta
  kernel:
    func: ftrl
    data_type: param

- op: fused_attention
  args: (Tensor x, Tensor ln_scale, Tensor ln_bias, Tensor qkv_weight, Tensor qkv_bias, Tensor cache_kv, Tensor src_mask, Tensor out_linear_weight, Tensor out_linear_bias, Tensor ln_scale_2, Tensor ln_bias_2, int num_heads, bool transpose_qkv_wb, bool pre_layer_norm, float epsilon, float attn_dropout_rate, bool is_test, bool attn_dropout_fix_seed, int attn_dropout_seed, str attn_dropout_implementation, float dropout_rate, bool dropout_fix_seed, int dropout_seed, str dropout_implementation, float ln_epsilon, bool add_residual, int ring_id)
  output: Tensor(ln_mean), Tensor(ln_var), Tensor(ln_out), Tensor(qkv_out), Tensor(qkv_bias_out), Tensor(transpose_out_2), Tensor(qk_out), Tensor(qktv_out), Tensor(softmax_out), Tensor(attn_dropout_mask_out), Tensor(attn_dropout_out), Tensor(src_mask_out), Tensor(fmha_out), Tensor(out_linear_out), Tensor(dropout_mask_out), Tensor(ln_mean_2), Tensor(ln_var_2), Tensor(bias_dropout_residual_out), Tensor(cache_kv_out), Tensor(out)
  kernel:
    func: fused_attention
    data_type : x
  infer_meta:
    func: FusedAttentionInferMeta
  optional: cache_kv, ln_scale, ln_bias, qkv_bias, src_mask, out_linear_bias, ln_scale_2, ln_bias_2, ln_mean_2, ln_var_2, bias_dropout_residual_out, cache_kv_out
  backward: fused_attention_grad

- op: fused_elemwise_add_activation
  args: (Tensor x, Tensor y, str[] functor_list, float scale=0.0, int axis=-1, bool save_intermediate_out=false)
  output: Tensor(out), Tensor(intermediate_out)
  kernel:
    func: fused_elemwise_add_activation
  infer_meta:
    func : FusedElemwiseAddActivationInferMeta
  backward: fused_elemwise_add_activation_grad
  intermediate: intermediate_out

- op: fused_feedforward
  args: (Tensor x, Tensor dropout1_seed, Tensor dropout2_seed, Tensor linear1_weight, Tensor linear1_bias, Tensor linear2_weight, Tensor linear2_bias, Tensor ln1_scale, Tensor ln1_bias, Tensor ln2_scale, Tensor ln2_bias, bool pre_layer_norm, float ln1_epsilon, float ln2_epsilon, str act_method, float dropout1_prob, float dropout2_prob, str dropout1_implementation, str dropout2_implementation, bool is_test, bool dropout1_fix_seed, bool dropout2_fix_seed, int dropout1_seed_val, int dropout2_seed_val, bool add_residual, int ring_id)
  output: Tensor(out), Tensor(dropout1_mask), Tensor(dropout2_mask), Tensor(ln1_mean), Tensor(ln1_variance), Tensor(ln2_mean), Tensor(ln2_variance), Tensor(linear1_out), Tensor(ln1_out), Tensor(dropout1_out), Tensor(dropout2_out)
  kernel:
    func: fused_feedforward
    data_type : x
  infer_meta:
    func: FusedFeedForwardInferMeta
  optional: dropout1_seed, dropout2_seed, linear1_bias, linear2_bias, ln1_scale, ln1_bias, ln2_scale, ln2_bias, ln2_mean, ln2_variance, ln1_mean, ln1_variance, ln1_out
  backward: fused_feedforward_grad

- op: lars_momentum
  args: (Tensor param,  Tensor velocity, Tensor grad, Tensor learning_rate, Tensor master_param, float mu, float lars_coeff=0.001f, float[] lars_weight_decay={0.0005}, float epsilon=0, bool multi_precision=false, float rescale_grad=1.0f)
  output: Tensor(param_out), Tensor(velocity_out), Tensor(master_param_out)
  infer_meta:
    func: SparseMomentumInferMeta
    param: [param, learning_rate, velocity]
  kernel:
    func: lars_momentum
    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
    data_type: param
  optional: master_param, master_param_out

- op: nce
  args: (Tensor input, Tensor label, Tensor weight, Tensor bias, Tensor sample_weight, Tensor custom_dist_probs, Tensor custom_dist_alias, Tensor custom_dist_alias_probs, int num_total_classes, int[] custom_neg_classes={}, int num_neg_samples=10, int sampler=0, int seed=0, bool is_sparse=false, bool remote_prefetch=false, bool is_test=false)
  output: Tensor(cost), Tensor(sample_logits), Tensor(sample_labels)
  infer_meta:
    func: NceInferMeta
  kernel:
    func: nce
    data_type: input
  optional: bias, sample_weight, custom_dist_probs, custom_dist_alias, custom_dist_alias_probs
  backward: nce_grad

- op: number_count
  args: (Tensor numbers, int upper_range)
  output: Tensor(out)
  infer_meta:
     func: NumberCountInferMeta
  kernel:
     func: number_count
     data_type: numbers

- op: onednn_to_paddle_layout
  args: (Tensor x, int dst_layout)
  output: Tensor(out)
  infer_meta:
    func : UnchangedInferMeta
    param : [x]
  kernel:
    func: onednn_to_paddle_layout

- op: sparse_momentum
  args: (Tensor param, Tensor grad, Tensor velocity, Tensor index, Tensor learning_rate, Tensor master_param,float mu, Scalar axis=0, bool use_nesterov=false,str regularization_method="", float regularization_coeff=0.0f, bool multi_precision=false, float rescale_grad=1.0f)
  output: Tensor(param_out), Tensor(velocity_out), Tensor(master_param_out)
  infer_meta:
    func: SparseMomentumInferMeta
    param: [param, learning_rate, velocity]
  kernel:
    func: sparse_momentum
    data_type: param
  optional: master_param, master_param_out
