- 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 : all_to_all args : (Tensor x, int ring_id = 0) output : Tensor(out) infer_meta : func : AllToAllInferMeta param: [x] kernel : func : all_to_all param: [x] - 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 : dist_concat args : (Tensor x, int ring_id = 0, int nranks = 1) output : Tensor(out) infer_meta : func : DistConcatInferMeta param: [x, nranks] kernel : func : dist_concat param: [x, nranks] - 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 backward : einsum_grad intermediate : inner_cache, xshape - op : elementwise_pow args : (Tensor x, Tensor y, int axis = -1) output : Tensor(out) infer_meta : func : ElementwiseRawInferMeta kernel : func : elementwise_pow backward : elementwise_pow_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 : exponential_ args : (Tensor x, float lam = 1.0f) output : Tensor(out) infer_meta : func : UnchangedInferMeta param : [x] kernel : func : exponential inplace : (x -> out) backward : exponential__grad - 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 : floor_divide args : (Tensor x, Tensor y, int axis = -1) output : Tensor(out) infer_meta : func : ElementwiseRawInferMeta kernel : func : floor_divide - 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 : full_like args : (Tensor x, Scalar value = 0.0, DataType dtype = DataType::UNDEFINED) output: Tensor(out) infer_meta : func : FillAnyLikeInferMeta kernel : func : full_like param : [x, value, dtype] data_type : dtype > x - 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 : matmul args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false) output : Tensor infer_meta : func : MatmulInferMeta kernel : func : matmul backward : matmul_grad - 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 : minimum args : (Tensor x, Tensor y, int axis = -1) output : Tensor(out) infer_meta : func : ElementwiseRawInferMeta kernel : func : minimum backward : minimum_grad - op : norm args : (Tensor x, int axis, float epsilon=1.0e-10f, bool is_test=false) output : Tensor(out), Tensor(norm) infer_meta : func : NormInferMeta kernel : func : norm backward : norm_grad intermediate : norm - 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", bool use_cudnn = false) 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", bool use_cudnn = false) 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 : prod args : (Tensor x, IntArray dims={0}, bool keep_dim=false, bool reduce_all=false, int in_dtype=-1, DataType out_dtype=DataType::UNDEFINED) output : Tensor(out) infer_meta : func : ReduceIntArrayAxisInferMetaBase param : [x, dims, keep_dim, reduce_all, out_dtype] kernel : func : prod param : [x, dims, keep_dim, reduce_all, out_dtype] data_type : x backward : prod_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 : remainder args : (Tensor x, Tensor y, int axis = -1) output : Tensor (out) infer_meta : func : ElementwiseRawInferMeta kernel : func : remainder inplace : (x -> out) - 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) output : Tensor(out) infer_meta : func : UnchangedInferMeta param : [x] kernel : func : swish 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