sparse_backward.yaml 14.3 KB
Newer Older
1
- backward_op : abs_grad
2 3 4
  forward : tanh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
5 6 7
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
8 9 10 11
  kernel :
    func : abs_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           abs_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

12
- backward_op : acos_grad
13 14 15
  forward : acos(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
16 17 18
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
19 20 21 22
  kernel :
    func : acos_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           acos_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

23
- backward_op : acosh_grad
24 25 26
  forward : acosh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
27 28 29
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
30 31 32 33
  kernel :
    func : acosh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           acosh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

34
- backward_op : add_grad
35 36 37
  forward : add(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
38
  infer_meta :
39
    func : GeneralBinaryGradInferMeta
40
    param : [x, y]
41 42
  kernel :
    func : add_coo_coo_grad{sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo},
43 44
           add_csr_csr_grad{sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr},
           add_coo_dense_grad{sparse_coo, dense, sparse_coo -> sparse_coo, dense}
45

46
- backward_op : addmm_grad
47 48 49
  forward : addmm(Tensor input, Tensor x, Tensor y, float alpha=1.0, float beta=1.0) -> Tensor(out)
  args : (Tensor input, Tensor x, Tensor y, Tensor out_grad, float alpha=1.0, float beta=1.0)
  output : Tensor(input_grad), Tensor(x_grad), Tensor(y_grad)
50 51 52
  infer_meta :
    func : GeneralTernaryGradInferMeta
    param : [input, x, y]
53 54 55 56 57 58
  kernel :
    func : addmm_csr_dense_grad {dense, sparse_csr, dense, dense -> dense, sparse_csr, dense},
           addmm_csr_csr_grad {sparse_csr, sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr, sparse_csr},
           addmm_coo_dense_grad {dense, sparse_coo, dense, dense -> dense, sparse_coo, dense},
           addmm_coo_coo_grad {sparse_coo, sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo, sparse_coo}

59
- backward_op : asin_grad
60 61 62
  forward : asin(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
63 64 65
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
66 67 68 69
  kernel :
    func : asin_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           asin_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

70
- backward_op : asinh_grad
71 72 73
  forward : asinh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
74 75 76
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
77 78 79 80
  kernel :
    func : asinh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           asinh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

81
- backward_op : atan_grad
82 83 84
  forward : atan(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
85 86 87
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
88 89 90 91
  kernel :
    func : atan_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           atan_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

92
- backward_op : atanh_grad
93 94 95
  forward : atanh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
96 97 98
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
99 100 101 102
  kernel :
    func : atanh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           atanh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

103
- backward_op : cast_grad
104 105 106
  forward : cast(Tensor x, DataType index_dtype, DataType value_dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad, DataType value_dtype)
  output : Tensor(x_grad)
107
  infer_meta :
108
    func : UnchangedInferMeta
109
    param: [x]
110 111 112 113 114
  kernel :
    func : cast_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           cast_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
    data_type : out_grad

115
- backward_op : conv3d_coo_grad
116 117
  forward : conv3d_coo (Tensor x, Tensor kernel, int[] paddings, int[] dilations, int[] strides, int groups, bool subm, str key) -> Tensor(out), Tensor(rulebook), Tensor(counter)
  args : (Tensor x, Tensor kernel, Tensor out, Tensor rulebook, Tensor counter, Tensor out_grad, int[] paddings, int[] dilations, int[] strides, int groups, bool subm, str key)
118
  output : Tensor(x_grad), Tensor(kernel_grad)
119 120 121
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, kernel]
122
  kernel :
123
    func : conv3d_coo_grad{sparse_coo, dense, sparse_coo, dense, dense, sparse_coo -> sparse_coo, dense}
Z
zhangkaihuo 已提交
124

125
- backward_op : divide_grad
126 127 128
  forward : divide(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
129
  infer_meta :
130
    func : GeneralBinaryGradInferMeta
131
    param : [x, y]
132 133 134 135
  kernel :
    func : divide_coo_coo_grad{sparse_coo, sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo},
           divide_csr_csr_grad{sparse_csr, sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr}

136
- backward_op : divide_scalar_grad
137 138 139 140 141
  forward : divide_scalar (Tensor x, float scalar) -> Tensor(out)
  args : (Tensor out_grad, float scalar)
  output : Tensor(x_grad)
  invoke : divide_scalar(out_grad, scalar)

142
- backward_op : expm1_grad
143 144 145
  forward : expm1(Tensor x) -> Tensor(out)
  args : (Tensor out, Tensor out_grad)
  output : Tensor(x_grad)
146 147 148
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
149 150 151 152
  kernel :
    func : expm1_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           expm1_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

153
- backward_op : leaky_relu_grad
154 155 156
  forward : leaky_relu(Tensor x, float alpha) -> Tensor(out)
  args : (Tensor x, Tensor out_grad, float alpha)
  output : Tensor(x_grad)
157 158 159
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
160 161 162 163
  kernel :
    func : leaky_relu_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           leaky_relu_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

164
- backward_op : log1p_grad
165 166 167
  forward : log1p(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
168 169 170
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
171 172 173 174
  kernel :
    func : log1p_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           log1p_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

175
- backward_op : masked_matmul_grad
176 177 178
  forward : masked_matmul(Tensor x, Tensor y, Tensor mask) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
179 180 181
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, y]
182
  kernel :
183
    func : masked_matmul_csr_grad{dense, dense, sparse_csr -> dense, dense}
184

185
- backward_op : matmul_grad
186 187 188
  forward : matmul(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
189 190 191
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, y]
192
  kernel :
193 194 195 196
    func : matmul_csr_dense_grad {sparse_csr, dense, dense -> sparse_csr, dense},
           matmul_csr_csr_grad {sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr},
           matmul_coo_dense_grad {sparse_coo, dense, dense -> sparse_coo, dense},
           matmul_coo_coo_grad {sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo}
197

198
- backward_op : maxpool_grad
199 200
  forward : maxpool(Tensor x, int[] kernel_sizes, int[] paddings, int[] dilations, int[] strides) -> Tensor(out), Tensor(rulebook), Tensor(counter)
  args : (Tensor x, Tensor rulebook, Tensor counter, Tensor out, Tensor out_grad, int[] kernel_sizes)
201
  output : Tensor(x_grad)
202 203 204
  infer_meta :
    func : UnchangedInferMeta
    param: [x]
205
  kernel :
206
    func : maxpool_coo_grad {sparse_coo, dense, dense, sparse_coo, sparse_coo -> sparse_coo}
207

208
- backward_op : multiply_grad
209 210 211
  forward : multiply(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
212
  infer_meta :
213
    func : GeneralBinaryGradInferMeta
214
    param : [x, y]
215 216 217 218
  kernel :
    func : multiply_coo_coo_grad{sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo},
           multiply_csr_csr_grad{sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr}

219
- backward_op : mv_grad
220 221 222
  forward : mv(Tensor x, Tensor vec) -> Tensor(out)
  args : (Tensor x, Tensor vec, Tensor out_grad)
  output : Tensor(x_grad), Tensor(vec_grad)
223 224 225
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, vec]
226 227 228 229
  kernel :
    func : mv_coo_grad{sparse_coo, dense, dense -> sparse_coo, dense},
           mv_csr_grad{sparse_csr, dense, dense -> sparse_csr, dense}

230
- backward_op : pow_grad
231 232 233
  forward : pow(Tensor x, float factor) -> Tensor(out)
  args : (Tensor x, Tensor out_grad, float factor)
  output : Tensor(x_grad)
234 235 236
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
237 238 239 240
  kernel :
    func : pow_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           pow_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

241
- backward_op : relu6_grad
242 243 244
  forward : relu6(Tensor x, float threshold) -> Tensor(out)
  args : (Tensor out, Tensor out_grad, float threshold)
  output : Tensor(x_grad)
245 246 247
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
248 249 250 251
  kernel :
    func : relu6_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           relu6_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

252
- backward_op : relu_grad
253
  forward : relu(Tensor x) -> Tensor(out)
254
  args : (Tensor out, Tensor out_grad)
255
  output : Tensor(x_grad)
256 257 258
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
Z
zhangkaihuo 已提交
259
  kernel :
260 261 262
    func : relu_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           relu_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

263
- backward_op : scale_grad
264 265 266
  forward : scale(Tensor x, float scale, float bias, bool bias_after_scale) -> Tensor(out)
  args : (Tensor out_grad, float scale)
  output : Tensor(x_grad)
267 268 269
  infer_meta :
    func : UnchangedInferMeta
    param : [out_grad]
270
  invoke : scale(out_grad, scale, 0.0, true)
Z
zhangkaihuo 已提交
271

272
- backward_op : sin_grad
273
  forward : sin(Tensor x) -> Tensor(out)
Z
zhangkaihuo 已提交
274
  args : (Tensor x, Tensor out_grad)
275
  output : Tensor(x_grad)
276 277 278
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
279
  kernel :
280 281 282
    func : sin_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sin_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

283
- backward_op : sinh_grad
284 285 286
  forward : sinh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
287 288 289
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
290 291 292
  kernel :
    func : sinh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sinh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
293

294
- backward_op : softmax_grad
295 296 297
  forward : softmax(Tensor x, int axis=-1) -> Tensor(out)
  args : (Tensor out, Tensor out_grad, int axis)
  output : Tensor(x_grad)
298 299 300
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
301 302 303
  kernel :
    func : softmax_csr_grad{sparse_csr, sparse_csr -> sparse_csr}

304
- backward_op : sparse_coo_tensor_grad
305 306 307
  forward : sparse_coo_tensor(Tensor values, Tensor indices, IntArray dense_shape) -> Tensor(out)
  args : (Tensor indices, Tensor out_grad)
  output : Tensor(values_grad)
308 309 310
  infer_meta :
    func : UnchangedInferMeta
    param: [out_grad]
311 312 313
  kernel :
    func : sparse_coo_tensor_grad{dense, sparse_coo -> dense}

314
- backward_op : sqrt_grad
315
  forward : sqrt(Tensor x) -> Tensor(out)
316
  args : (Tensor out, Tensor out_grad)
317
  output : Tensor(x_grad)
318 319 320
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
321
  kernel :
322 323 324
    func : sqrt_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sqrt_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

325
- backward_op : square_grad
326 327 328
  forward : square(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
329 330 331
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
332 333 334
  kernel :
    func : square_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           square_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
335

336
- backward_op : subtract_grad
337 338 339
  forward : subtract(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
340
  infer_meta :
341
    func : GeneralBinaryGradInferMeta
342
    param : [x, y]
343 344 345 346
  kernel :
    func : subtract_coo_coo_grad{sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo},
           subtract_csr_csr_grad{sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr}

347
- backward_op : tan_grad
348 349 350
  forward : tan(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
351 352 353
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
354 355 356 357
  kernel :
    func : tan_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           tan_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

358
- backward_op : tanh_grad
359
  forward : tanh(Tensor x) -> Tensor(out)
360
  args : (Tensor out, Tensor out_grad)
361
  output : Tensor(x_grad)
362 363 364
  infer_meta :
    func : UnchangedInferMeta
    param : [out]
365
  kernel :
366 367
    func : tanh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           tanh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
368

369
- backward_op : to_dense_grad
370 371 372
  forward : to_dense(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
373 374 375
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
376 377 378
  kernel :
    func : coo_to_dense_grad{sparse_coo, dense -> sparse_coo}

379
- backward_op : to_sparse_coo_grad
380 381 382
  forward : to_sparse_coo(Tensor x, int64_t sparse_dim) -> Tensor(out)
  args : (Tensor out_grad)
  output : Tensor(x_grad)
383 384
  infer_meta :
    func : UnchangedInferMeta
385 386 387
  kernel :
    func : coo_to_dense { sparse_coo -> dense }

388
- backward_op : values_grad
389
  forward : values_coo(Tensor x) -> Tensor(out)
390 391
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
392 393 394
  infer_meta :
    func : UnchangedInferMeta
    param : [x]
395
  kernel :
396
    func : values_coo_grad{sparse_coo, dense-> sparse_coo}
397

398
- backward_op: fused_attention_grad
399 400
  forward : fused_attention_csr(Tensor query, Tensor key, Tensor value, Tensor sparse_mask, Tensor key_padding_mask, Tensor attn_mask) -> Tensor(out), Tensor(softmax)
  args: (Tensor query, Tensor key, Tensor value, Tensor softmax, Tensor out_grad)
401
  output : Tensor(query_grad), Tensor(key_grad), Tensor(value_grad)
402
  infer_meta :
403
    func : sparse::FusedAttentionGradInferMeta
404 405
  kernel :
    func : fused_attention_csr_grad{dense, dense, dense, sparse_csr, dense -> dense, dense, dense}
406 407
    layout : softmax
    data_type: query