sparse_backward.yaml 14.2 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)
Z
zhangkaihuo 已提交
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)
Z
zhangkaihuo 已提交
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)
Z
zhangkaihuo 已提交
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)
Z
zhangkaihuo 已提交
38 39 40
  infer_meta :
    func : GeneralBinaryGradInferMeta 
    param : [x, y]
41 42 43 44
  kernel :
    func : add_coo_coo_grad{sparse_coo, sparse_coo, sparse_coo -> sparse_coo, sparse_coo},
           add_csr_csr_grad{sparse_csr, sparse_csr, sparse_csr -> sparse_csr, sparse_csr}

45
- backward_op : addmm_grad
46 47 48
  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)
Z
zhangkaihuo 已提交
49 50 51
  infer_meta :
    func : GeneralTernaryGradInferMeta
    param : [input, x, y]
52 53 54 55 56 57
  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}

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

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

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

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

102
- backward_op : cast_grad
103 104 105
  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)
Z
zhangkaihuo 已提交
106 107 108
  infer_meta :
    func : UnchangedInferMeta 
    param: [x]
109 110 111 112 113
  kernel :
    func : cast_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           cast_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
    data_type : out_grad

114
- backward_op : conv3d_coo_grad
115 116
  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)
117
  output : Tensor(x_grad), Tensor(kernel_grad)
Z
zhangkaihuo 已提交
118 119 120
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, kernel]
121
  kernel :
122
    func : conv3d_coo_grad{sparse_coo, dense, sparse_coo, dense, dense, sparse_coo -> sparse_coo, dense}
Z
zhangkaihuo 已提交
123

124
- backward_op : divide_grad
125 126 127
  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)
Z
zhangkaihuo 已提交
128 129 130
  infer_meta :
    func : GeneralBinaryGradInferMeta 
    param : [x, y]
131 132 133 134
  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}

135
- backward_op : divide_scalar_grad
136 137 138 139 140
  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)

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

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

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

174
- backward_op : masked_matmul_grad
175 176 177
  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)
Z
zhangkaihuo 已提交
178 179 180
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, y]
181
  kernel :
182
    func : masked_matmul_csr_grad{dense, dense, sparse_csr -> dense, dense}
183

184
- backward_op : matmul_grad
185 186 187
  forward : matmul(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
Z
zhangkaihuo 已提交
188 189 190
  infer_meta :
    func : GeneralBinaryGradInferMeta
    param : [x, y]
191
  kernel :
192 193 194 195
    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}
196

197
- backward_op : maxpool_grad
198 199
  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)
200
  output : Tensor(x_grad)
Z
zhangkaihuo 已提交
201 202 203
  infer_meta :
    func : UnchangedInferMeta
    param: [x]
204
  kernel :
205
    func : maxpool_coo_grad {sparse_coo, dense, dense, sparse_coo, sparse_coo -> sparse_coo}
206

207
- backward_op : multiply_grad
208 209 210
  forward : multiply(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
Z
zhangkaihuo 已提交
211 212 213
  infer_meta :
    func : GeneralBinaryGradInferMeta 
    param : [x, y]
214 215 216 217
  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}

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

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

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

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

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

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

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

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

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

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

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

335
- backward_op : subtract_grad
336 337 338
  forward : subtract(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
Z
zhangkaihuo 已提交
339 340 341
  infer_meta :
    func : GeneralBinaryGradInferMeta 
    param : [x, y]
342 343 344 345
  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}

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

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

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

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

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

397
- backward_op: fused_attention_grad
398 399
  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)
400
  output : Tensor(query_grad), Tensor(key_grad), Tensor(value_grad)
Z
zhangkaihuo 已提交
401 402
  infer_meta :
    func : sparse::FusedAttentionGradInferMeta 
403 404
  kernel :
    func : fused_attention_csr_grad{dense, dense, dense, sparse_csr, dense -> dense, dense, dense}
405 406
    layout : softmax
    data_type: query