static_ops.yaml

- op : all
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMetaBase
  kernel :
    func : all

- op : all_gather
  args : (Tensor x, int ring_id = 0, int nranks=0)
  output : Tensor(out)
  infer_meta :
    func : AllGatherInferMeta
    param: [x, nranks]
  kernel :
    func : all_gather
    param: [x, nranks]

- op : all_reduce
  args : (Tensor x, int ring_id = 0, int reduce_type = 0)
  output : Tensor(out)
  infer_meta :
    func : AllReduceInferMeta
    param: [x]
  kernel :
    func : all_reduce
    param: [x, reduce_type]

- op : amax
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
    param : [x, axis, keepdim]
  kernel :
    func : amax_raw
    param : [x, axis, keepdim, reduce_all]
  backward : amax_grad

- op : amin
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMeta
    param : [x, axis, keepdim]
  kernel :
    func : amin_raw
    param : [x, axis, keepdim, reduce_all]
  backward : amin_grad

- op : any
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMetaBase
    param : [x, axis, keepdim, reduce_all]
  kernel :
    func : any_raw
    param : [x, axis, keepdim, reduce_all]

- op : arange
  args : (Tensor start, Tensor end, Tensor step)
  output : Tensor(out)
  infer_meta :
    func : ArangeInferMeta
  kernel :
    func : arange
  data_transform :
    skip_transform : start, end, step

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

- op : assign_value
  args : (int[] shape, DataType dtype, int[] bool_values = {}, float[] fp32_values = {}, int[] int32_values = {}, int64_t[] int64_values = {})
  output : Tensor(out)
  infer_meta :
    func : AssignValueInferMeta
    param : [shape, dtype]
  kernel :
    func : assign_value
    param : [shape, dtype, values]
    data_type : dtype

- op : broadcast
  args : (Tensor x, int ring_id = 0, int root = 0)
  output : Tensor(out)
  infer_meta :
    func : DistBroadcastInferMeta
    param: [x]
  kernel :
    func : broadcast
    param: [x, root]

- op : conv2d_transpose
  args : (Tensor x, Tensor filter, Tensor bias, 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
    param : [x, filter, strides, paddings, output_padding, output_size, padding_algorithm, groups, dilations, data_format]
  kernel :
    func : conv2d_transpose
    param : [x, filter, strides, paddings, output_padding, output_size, padding_algorithm, groups, dilations, data_format]
    data_type : x
  optional : bias
  backward : conv2d_transpose_grad

- op : decode_jpeg
  args : (Tensor x, str mode = "unchanged")
  output : Tensor(out)
  infer_meta :
    func : DecodeJpegInferMeta
    param : [x, mode]
  kernel :
    func : decode_jpeg
    param : [x, mode]

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

- op : depthwise_conv2d_transpose
  args : (Tensor x, Tensor filter, Tensor bias, 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
    param : [x, filter, strides, paddings, output_padding, output_size, padding_algorithm, groups, dilations, data_format]
  kernel :
    func : depthwise_conv2d_transpose
    param : [x, filter, strides, paddings, output_padding, output_size, padding_algorithm, groups, dilations, data_format]
    data_type : x
  optional : bias
  backward : depthwise_conv2d_transpose_grad

- op : embedding
  args : (Tensor x, Tensor weight, int64_t padding_idx=-1)
  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)
  output: Tensor(out)
  infer_meta :
    func : CreateInferMeta
    param : [shape, dtype]
  kernel :
    func : empty
    param : [shape, dtype]
    data_type : dtype

- op : equal
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : equal_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

- op : eye
  args : (Scalar(int64_t) num_rows, Scalar(int64_t) num_columns = -1, DataType dtype = DataType::FLOAT32)
  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

- op : frobenius_norm
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceInferMetaBase
  kernel :
    func : frobenius_norm
    param : [x, axis, keepdim, reduce_all]
  backward : frobenius_norm_grad

- op : gaussian
  args : (IntArray shape = {}, float mean = .0f, float std = 1.0f, int seed = 0, DataType dtype = DataType::FLOAT32)
  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

- op : greater_equal
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : greater_equal_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

- op : greater_than
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : greater_than_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

- op : hardswish
  args : (Tensor x, float threshold = 6.0f, float scale = 6.0f, float offset = 3.0f)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : hardswish
    param : [x]
  backward : hardswish_grad

- op : less_equal
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : less_equal_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

- op : less_than
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : less_than_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

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

- op : matrix_rank
  args : (Tensor x, Tensor tol_tensor, float tol=0.0f, bool hermitian=false, bool use_default_tol=true)
  output : Tensor(out)
  infer_meta :
    func : MatrixRankStaticInferMeta
    param : [x, tol_tensor, use_default_tol, hermitian]
  optional : tol_tensor
  kernel :
    func : matrix_rank {dense -> dense},
           matrix_rank_tol {dense, dense -> dense}
    data_type : x

- op : max
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMetaBase
    param : [x, axis, keepdim, reduce_all]
  kernel :
    func : max_raw
    param : [x, axis, keepdim, reduce_all]
  backward : max_grad

- op : maximum
  args : (Tensor x, Tensor y, int axis = -1)
  output : Tensor(out)
  infer_meta :
    func : ElementwiseRawInferMeta
  kernel :
    func : maximum
  backward : maximum_grad

- op : min
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
  output : Tensor(out)
  infer_meta :
    func : ReduceIntArrayAxisInferMetaBase
    param : [x, axis, keepdim, reduce_all]
  kernel :
    func : min_raw
    param : [x, axis, keepdim, reduce_all]
  backward : min_grad

- op : not_equal
  args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
  output : Tensor(out)
  infer_meta :
    func : CompareRawInferMeta
    param : [x, y, axis]
  kernel :
    func : not_equal_raw
    param : [x, y, axis]
    backend : x
    force_backend : force_cpu

- op : one_hot
  args : (Tensor x, Scalar(int) depth = -1, DataType dtype = DataType::FLOAT32, bool allow_out_of_range = false)
  output : Tensor(out)
  infer_meta :
    func : OneHotRawInferMeta
  kernel :
    func : one_hot_raw
    data_type : x

- op : p_recv
  args : (int ring_id = 0, int peer = 0, DataType dtype = DataType::FLOAT32, bool dynamic_shape = false)
  output : Tensor(out)
  infer_meta :
    func : PRecvInferMeta
    param : [peer, dtype]
  kernel :
    func : p_recv
    param : [peer, dtype, dynamic_shape]
    data_type : dtype

- op : p_recv_array
  args : (int ring_id = 0, int peer = 0, DataType dtype = DataType::FLOAT32, int[] out_shape = {})
  output : Tensor(out)
  infer_meta :
    func : PRecvArrayInferMeta
    param : [peer, dtype, out_shape]
  kernel :
    func : p_recv_array
    param : [peer, dtype, out_shape]

- op : pool2d
  args : (Tensor x, IntArray kernel_size, int[] strides = {1,1}, int[] paddings = {0,0}, bool ceil_mode = false, bool exclusive = true, str data_format = "NCHW", str pooling_type = "", bool global_pooling = false, bool adaptive = false, str padding_algorithm = "EXPLICIT")
  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 = {1,1,1}, int[] paddings = {0,0,0}, bool ceil_mode = false, bool exclusive = true, str data_format = "NCDHW", str pooling_type = "", bool global_pooling = false, bool adaptive = false, str padding_algorithm = "EXPLICIT")
  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 : randint
  args : (int low, int high, IntArray shape = {}, DataType dtype = DataType::INT64, int seed = 0)
  output : Tensor(out)
  infer_meta :
    func : RandintInferMeta
    param : [low, high, shape, dtype]
  kernel :
    func : randint
    param : [low, high, shape, dtype]
    data_type : dtype

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

- op : reduce
  args : (Tensor x, int ring_id = 0, int root_id = 0, int reduce_type = 0)
  output : Tensor(out)
  infer_meta :
    func : DistReduceInferMeta
    param: [x]
  kernel :
    func : reduce
    param: [x, root_id, reduce_type]

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

- op : relu6
  args : (Tensor x, float threshold = 6.0f)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : relu6_raw
  backward : relu6_grad

- op : rnn
  args: (Tensor x, Tensor[] pre_state, Tensor[] weight_list, Tensor sequence_length, 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, dropout_state_out
  intermediate : reserve

- op : share_buffer
  args : (Tensor[] x, bool[] share_dims_and_dtype={})
  output : Tensor[](out){x.size()}, Tensor[](xout){x.size()}
  infer_meta :
    func : ShareBufferInferMeta
  kernel :
    func : share_buffer

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

- op : strided_slice
  args : (Tensor x, int[] axes, IntArray starts={}, IntArray ends={}, IntArray strides={}, int[] infer_flags={}, int[] decrease_axis={})
  output : Tensor
  infer_meta :
    func : StridedSliceRawInferMeta
  kernel :
    func : strided_slice
    param : [x, axes, starts, ends, strides]
  backward : strided_slice_grad

- op : sum
  args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, DataType out_dtype=DataType::UNDEFINED)
  output : Tensor(out)
  infer_meta :
    func : SumRawInferMeta
    param : [x, axis, keepdim, reduce_all, out_dtype]
  kernel :
    func : sum_raw
    param : [x, axis, keepdim, reduce_all, out_dtype]
    data_type : x
  backward : sum_grad

- op : swish
  args : (Tensor x, float beta = 1.0f)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : swish_raw
  backward : swish_grad

- op : tril_indices
  args : (int rows = 0, int cols = 0, int offset = 0, DataType dtype = DataType::INT64)
  output : Tensor(out)
  infer_meta :
    func : TrilIndicesInferMeta
    param : [rows, cols, offset, dtype]
  kernel :
    func : tril_indices
    param : [rows, cols, offset, dtype]
    data_type : dtype

- op : tril_triu
  args : (Tensor x, int diagonal = 0, bool lower = false)
  output : Tensor(out)
  infer_meta :
    func : TrilTriuInferMeta
  kernel :
    func : tril_triu
  backward : tril_triu_grad

- op : triu_indices
  args : (int row = 0, int col = 0, int offset = 0, DataType dtype = DataType::INT64)
  output : Tensor(out)
  infer_meta :
    func : TriuIndicesInferMeta
    param : [row, col, offset, dtype]
  kernel :
    func : triu_indices
    param : [row, col, offset, dtype]
    data_type : dtype

- op : truncated_gaussian_random
  args : (int[] shape, float mean = .0f, float std = 1.0f, int seed = 0, DataType dtype=DataType::FLOAT32)
  output : Tensor(out)
  infer_meta :
    func : TruncatedGaussianRandomInferMeta
    param : [shape, mean, std, seed, dtype]
  kernel :
    func : truncated_gaussian_random
    param : [shape, mean, std, seed, dtype]
    data_type : dtype

- op : uniform
  args : (IntArray shape = {},  DataType dtype = DataType::FLOAT32, Scalar min = -1.0f, Scalar max = 1.0f, int seed = 0, int diag_num = 0, int diag_step = 0, float diag_val = 1.0f)
  output : Tensor(out)
  infer_meta :
    func : UniformRandomInferMeta
    param: [shape, dtype]
  kernel :
    func : uniform
    param: [shape, dtype, min, max, seed]
    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

- op : unpool
  args: (Tensor x, Tensor indices, int[] ksize, str unpooling_type, int[] strides = {1,1}, int[] paddings ={0,0} ,IntArray output_size = {0,0}, str data_format="NCHW")
  output: Tensor(out)
  infer_meta:
    func: UnpoolInferMeta
    param : [x, indices, ksize, strides, paddings,output_size, data_format]
  kernel:
    func: unpool
    data_type: x
    param : [x, indices, ksize, strides, paddings,output_size, data_format]
  backward: unpool_grad