static_ops.yaml 17.3 KB
Newer Older
1 2 3 4 5 6 7 8
- 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

9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
- 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]

29 30 31 32 33 34 35 36 37 38 39
- 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

40 41 42 43 44 45 46 47 48 49 50
- 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

51 52 53 54 55 56 57 58 59 60
- 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]

61 62 63 64 65 66 67 68 69 70
- 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

71 72 73 74 75 76 77 78 79 80 81
- op : assign
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : assign
  optional : x
  inplace : (x -> out)
  backward : assign_grad

82 83 84 85 86 87 88 89 90 91 92
- 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

93 94 95 96 97 98 99 100 101 102
- 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]

103 104 105 106 107 108 109 110 111 112 113 114 115
- 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

116 117 118 119 120 121 122 123 124 125
- 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]

126 127 128 129 130 131 132 133 134 135
- 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

136 137 138 139 140 141 142 143 144 145 146 147 148
- 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

149 150 151 152 153 154 155 156 157 158 159 160 161
- 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

162 163 164 165 166 167 168 169 170 171 172
- 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

173 174 175 176 177 178 179 180 181 182 183 184
- 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

185 186 187 188 189 190 191 192 193 194 195
- 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

196 197 198 199 200 201 202 203 204 205
- 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

206 207 208 209 210 211 212 213 214 215 216
- 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

217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
- 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

241 242 243 244 245 246 247 248 249 250 251
- 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

252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275
- 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

276 277 278 279 280 281 282 283 284 285 286 287 288
- 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

289 290 291 292 293
- 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
Z
zhangyuqin1998 已提交
294
    param : [x, tol_tensor, use_default_tol, hermitian]
295 296 297 298 299 300
  optional : tol_tensor
  kernel :
    func : matrix_rank {dense -> dense},
           matrix_rank_tol {dense, dense -> dense}
    data_type : x

301 302 303 304 305 306 307 308 309 310 311
- 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

312 313 314 315 316 317 318 319 320 321 322
- 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

323 324 325 326 327 328 329 330 331 332 333 334
- 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

335 336 337 338 339 340 341 342 343
- 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

344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364
- 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]

365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
- 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

387 388 389 390 391 392 393 394 395 396 397
- 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

398 399 400 401 402 403 404 405 406 407 408
- 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

409 410 411 412 413 414 415 416 417 418
- 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]

419 420 421 422 423 424 425 426 427 428
- 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]

429 430 431 432 433 434 435 436 437 438
- op : relu6
  args : (Tensor x, float threshold = 6.0f)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : relu6_raw
  backward : relu6_grad

439 440 441 442 443 444 445 446 447 448 449 450 451 452
- 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

453 454 455 456 457 458 459
- 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
460

461 462 463 464 465 466 467 468 469 470
- op : softmax
  args : (Tensor x, int axis = -1)
  output : Tensor(out)
  infer_meta :
    func : SoftmaxInferMeta
  kernel :
    func : softmax
  inplace : (x -> out)
  backward : softmax_grad

471 472 473 474 475 476 477 478 479 480
- 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

481 482 483 484 485 486 487 488 489 490 491 492
- 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

493 494 495 496 497 498 499 500 501 502
- 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

503 504 505 506 507 508 509 510 511 512 513
- 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

514 515 516 517 518 519 520 521 522
- 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

523 524 525 526 527 528 529 530 531 532 533
- 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

534 535 536 537 538 539 540 541 542 543 544
- 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

545 546 547 548 549 550 551 552 553 554
- 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
L
lzydev 已提交
555 556 557 558 559 560 561 562 563 564

- 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
565 566 567 568 569 570 571 572 573 574 575 576

- 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