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 : 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 : 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 : 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 : 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 : 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 : 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 : 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 : 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 : 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 : 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