ops.yaml 8.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
- op : acos
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : acos
  backward : acos_grad

- op : acosh
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : acosh
  backward : acosh_grad

19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36
- op : angle
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : RealAndImagInferMeta
  kernel :
    func : angle
  backward : angle_grad

- op : argsort
  args : (Tensor x, int axis=-1, bool descending=false)
  output : Tensor(out), Tensor(indices)
  infer_meta :
    func : ArgsortInferMeta
  kernel :
    func : argsort
  backward : argsort_grad

37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
- op : asin
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : asin
  backward : asin_grad

- op : asinh
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : asinh
  backward : asinh_grad

- op : atan
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : atan
  backward : atan_grad

64
- op : atan2
65 66 67 68 69 70 71 72
  args : (Tensor x, Tensor y)
  output : Tensor
  infer_meta :
    func : Atan2InferMeta
  kernel :
    func : atan2
  backward : atan2_grad

73 74 75 76 77 78 79 80 81
- op : atanh
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : atanh
  backward : atanh_grad

82
- op : bernoulli
83
  args : (Tensor x)
84
  output : Tensor(out)
85 86 87 88 89
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : bernoulli

90 91 92 93 94 95 96 97 98
- op : bmm
  args : (Tensor x, Tensor y)
  output : Tensor
  infer_meta :
    func : BmmInferMeta
  kernel :
    func : bmm
  backward : bmm_grad

99 100 101 102 103 104 105 106 107 108
- op : ceil
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : ceil
  inplace : (x -> out)
  backward : ceil_grad

109
- op : cholesky
110 111 112 113 114 115 116 117
  args : (Tensor x, bool upper=false)
  output : Tensor
  infer_meta :
    func : CholeskyInferMeta
  kernel :
    func : cholesky
  backward : cholesky_grad

118
- op : cholesky_solve
119 120 121 122 123 124 125 126
  args : (Tensor x, Tensor y, bool upper=false)
  output : Tensor
  infer_meta :
    func : CholeskySolveInferMeta
  kernel :
    func : cholesky_solve
  backward : cholesky_solve_grad

127
- op : cos
128
  args : (Tensor x)
129
  output : Tensor
130 131 132
  infer_meta :
    func : UnchangedInferMeta
  kernel :
133 134 135 136 137 138 139 140 141 142 143
    func : cos
  backward : cos_grad

- op : cosh
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : cosh
  backward : cosh_grad
144

145
- op : cross
146 147 148 149 150 151 152 153 154
  args : (Tensor x, Tensor y, int axis = 9)
  output : Tensor
  infer_meta :
    func : CrossInferMeta
  kernel :
    func : cross
    data_type : x
  backward : cross_grad

155 156 157 158 159 160 161 162 163
- op : det
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : determinant
  backward : det_grad

164
- op : diag
165 166 167 168 169 170 171 172
  args : (Tensor x, int offset = 0, float padding_value = 0.0)
  output : Tensor
  infer_meta :
    func : DiagInferMeta
  kernel :
    func : diag
  backward : diag_grad

173
- op : diagonal
174 175 176 177 178 179 180 181
  args : (Tensor x, int offset = 0, int axis1 = 0, int axis2 = 1)
  output : Tensor
  infer_meta :
    func : DiagonalInferMeta
  kernel :
    func : diagonal
  backward : diagonal_grad

182
- op : digamma
183 184 185 186 187 188 189 190
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : digamma
  backward : digamma_grad

191
- op : dist
192 193 194 195 196 197 198 199
  args : (Tensor x, Tensor y, float p = 2.0)
  output : Tensor
  infer_meta :
    func : DistInferMeta
  kernel :
    func : dist
  backward : dist_grad

200
- op : dot
201 202 203 204 205 206 207 208 209
  args : (Tensor x, Tensor y)
  output : Tensor
  infer_meta :
    func : DotInferMeta
  kernel :
    func : dot
    data_type : x
  backward : dot_grad

210
- op : erf
211 212 213 214 215 216 217 218
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : erf
  backward : erf_grad

219
- op : erfinv
220 221 222 223 224 225 226 227 228
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : erfinv
  inplace : (x -> out)
  backward : erfinv_grad

229 230 231 232 233 234 235 236 237 238
- op : exp
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : exp
  inplace : (x -> out)
  backward : exp_grad

239 240 241 242 243 244 245 246 247 248
- op : expm1
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : expm1
  backward : expm1_grad

249
- op : fft_c2c
F
Feiyu Chan 已提交
250 251 252 253 254 255 256 257
  args : (Tensor x, int64_t[] axes, str normalization, bool forward)
  output : Tensor
  infer_meta :
    func : FFTC2CInferMeta
  kernel :
    func : fft_c2c
  backward : fft_c2c_grad

258
- op : fft_c2r
F
Feiyu Chan 已提交
259 260 261 262 263 264 265 266
  args : (Tensor x, int64_t[] axes, str normalization, bool forward, int64_t last_dim_size=0L)
  output : Tensor
  infer_meta :
    func : FFTC2RInferMeta
  kernel :
    func : fft_c2r
  backward : fft_c2r_grad

267
- op : fft_r2c
F
Feiyu Chan 已提交
268 269 270 271 272 273 274 275
  args : (Tensor x, int64_t[] axes, str normalization, bool forward, bool onesided)
  output : Tensor
  infer_meta :
    func : FFTR2CInferMeta
  kernel :
    func : fft_r2c
  backward : fft_r2c_grad

276 277 278 279 280 281 282 283 284
- op : flip
  args : (Tensor x, int[] axis)
  output : Tensor (out)
  infer_meta :
    func : FlipInferMeta
  kernel :
    func : flip
  backward : flip_grad

285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314
- op : floor
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : floor
  inplace : (x -> out)
  backward : floor_grad

- op : hardshrink
  args : (Tensor x, float threshold = 0.5)
  output : Tensor (out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : hard_shrink
  backward : hardshrink_grad

- op : hardsigmoid
  args : (Tensor x, float slope = 0.2, float offset = 0.5)
  output : Tensor (out)
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : hard_sigmoid
  backward : hardsigmoid_grad

315
- op : lgamma
316 317 318 319 320 321
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : lgamma
F
Feiyu Chan 已提交
322
  backward : lgamma_grad
323

324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369
- op : log10
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : log10
  backward: log10_grad

- op : log1p
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : log1p
  backward: log1p_grad

- op : log2
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : log2
  backward: log2_grad

- op : logit
  args : (Tensor x, float eps = 1e-6f)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
  kernel :
    func : logit
  backward : logit_grad

- op : logsigmoid
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : logsigmoid
  backward : logsigmoid_grad

370
- op : mv
371 372 373 374 375 376 377 378
  args : (Tensor x, Tensor vec)
  output : Tensor
  infer_meta :
    func : MvInferMeta
  kernel :
    func : mv
  backward : mv_grad

379
- op : poisson
380 381 382 383 384 385 386 387
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : poisson
  backward : poisson_grad

388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407
- op : reciprocal
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : reciprocal
  inplace : (x -> out)
  backward : reciprocal_grad

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

408 409 410 411 412 413 414 415 416 417
- op : send_uv
  args : (Tensor x, Tensor y, Tensor src_index, Tensor dst_index, str message_op = "ADD")
  output : Tensor(out)
  infer_meta :
    func : SendUVInferMeta
  kernel :
    func : send_uv
    data_type : x
  backward : send_uv_grad

418 419 420 421 422 423 424 425 426
- op : silu
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : silu
  backward : silu_grad

427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444
- op : sin
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : sin
  backward : sin_grad

- op : sinh
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : sinh
  backward : sinh_grad

445
- op : solve
446 447 448 449 450 451 452 453 454
  args : (Tensor x, Tensor y)
  output : Tensor
  infer_meta :
    func : SolveInferMeta
  kernel :
    func : solve
    data_type : x
  backward : solve_grad

455 456 457 458 459 460 461 462 463
- op : tan
  args : (Tensor x)
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : tan
  backward : tan_grad

464 465 466 467 468 469 470 471 472 473
- op : tanh
  args : (Tensor x)
  output : Tensor(out)
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : tanh
  inplace : (x -> out)
  backward : tanh_grad

474
- op : trace
475 476 477 478 479 480 481 482
  args : (Tensor x, int offset = 0, int axis1 = 0, int axis2 = 1)
  output : Tensor
  infer_meta :
    func : TraceInferMeta
  kernel :
    func : trace
  backward : trace_grad

483
- op : trunc
484
  args : (Tensor input)
485 486 487 488 489 490
  output : Tensor
  infer_meta :
    func : UnchangedInferMeta
  kernel :
    func : trunc
  backward : trunc_grad