sparse_bw_api.yaml 10.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
- backward_api : abs_grad
  forward : tanh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : abs_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           abs_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : acos_grad
  forward : acos(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : acos_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           acos_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : acosh_grad
  forward : acosh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : acosh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           acosh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

25 26 27 28 29 30 31 32
- backward_api : add_grad
  forward : add(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
  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}

33 34 35 36 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 64 65 66 67 68 69 70 71 72 73
- backward_api : asin_grad
  forward : asin(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : asin_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           asin_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : asinh_grad
  forward : asinh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : asinh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           asinh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : atan_grad
  forward : atan(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : atan_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           atan_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : atanh_grad
  forward : atanh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : atanh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           atanh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : cast_grad
  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)
  kernel :
    func : cast_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           cast_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
    data_type : out_grad

74
- backward_api : conv3d_grad
Z
zhangkaihuo 已提交
75 76
  forward : conv3d (Tensor x, Tensor kernel, int[] paddings, int[] dilations, int[] strides, int groups, bool subm) -> Tensor(out@SparseCooTensor), Tensor(rulebook@DenseTensor)
  args : (Tensor x, Tensor kernel, Tensor rulebook, Tensor out_grad, int[] paddings, int[] dilations, int[] strides, int groups, bool subm)
77
  output : Tensor(x_grad), Tensor(kernel_grad)
78
  kernel :
Z
zhangkaihuo 已提交
79
    func : conv3d_coo_grad{sparse_coo, dense, dense, sparse_coo -> sparse_coo, dense}
Z
zhangkaihuo 已提交
80

81
- backward_api : coo_to_dense_grad
82
  forward : coo_to_dense(Tensor x) -> Tensor(out)
83
  args : (Tensor x, Tensor out_grad)
84
  output : Tensor(x_grad)
85
  kernel :
86
    func : sparse_coo_to_dense_grad{sparse_coo, dense-> sparse_coo}
87

88
- backward_api : create_sparse_coo_tensor_grad
89
  forward : create_sparse_coo_tensor(Tensor values, Tensor indices, IntArray dense_shape) -> Tensor(out)
90
  args : (Tensor indices, Tensor out_grad)
91
  output : Tensor(values_grad)
92
  kernel :
93
    func : sparse_coo_tensor_grad{dense, sparse_coo -> dense}
94

95
- backward_api : dense_to_coo_grad
96
  forward : dense_to_coo(Tensor x, int64_t sparse_dim) -> Tensor(out)
97
  args : (Tensor out_grad)
98
  output : Tensor(x_grad)
99 100
  invoke : to_dense_impl(out_grad)

101 102 103 104 105 106 107 108
- backward_api : divide_grad
  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)
  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}

109 110 111 112 113 114 115 116 117 118 119 120 121 122
- backward_api : divide_scalar_grad
  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)

- backward_api : log1p_grad
  forward : log1p(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : log1p_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           log1p_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

123 124 125 126 127
- backward_api : masked_matmul_grad
  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)
  kernel :
128
    func : masked_matmul_csr_grad{dense, dense, sparse_csr -> dense, dense}
129 130 131 132 133 134

- backward_api : matmul_grad
  forward : matmul(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
  kernel :
135 136 137 138
    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}
139

140 141 142 143 144 145 146
- backward_api : maxpool_grad
  forward : maxpool(Tensor x, int[] kernel_sizes, int[] paddings, int[] dilations, int[] strides) -> Tensor(out), Tensor(rulebook)
  args : (Tensor x, Tensor rulebook, Tensor out, Tensor out_grad, int[] kernel_sizes)
  output : Tensor(x_grad)
  kernel :
    func : maxpool_coo_grad {sparse_coo, dense, sparse_coo, sparse_coo -> sparse_coo}

147 148 149 150 151 152 153 154
- backward_api : multiply_grad
  forward : multiply(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
  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}

155 156 157 158 159 160 161 162
- backward_api : mv_grad
  forward : mv(Tensor x, Tensor vec) -> Tensor(out)
  args : (Tensor x, Tensor vec, Tensor out_grad)
  output : Tensor(x_grad), Tensor(vec_grad)
  kernel :
    func : mv_coo_grad{sparse_coo, dense, dense -> sparse_coo, dense},
           mv_csr_grad{sparse_csr, dense, dense -> sparse_csr, dense}

163 164 165 166 167 168 169 170
- backward_api : pow_grad
  forward : pow(Tensor x, float factor) -> Tensor(out)
  args : (Tensor x, Tensor out_grad, float factor)
  output : Tensor(x_grad)
  kernel :
    func : pow_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           pow_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

171 172
- backward_api : relu_grad
  forward : relu(Tensor x) -> Tensor(out)
173
  args : (Tensor out, Tensor out_grad)
174
  output : Tensor(x_grad)
Z
zhangkaihuo 已提交
175
  kernel :
176 177 178 179 180 181 182 183
    func : relu_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           relu_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : scale_grad
  forward : scale(Tensor x, float scale, float bias, bool bias_after_scale) -> Tensor(out)
  args : (Tensor out_grad, float scale)
  output : Tensor(x_grad)
  invoke : scale(out_grad, scale, 0.0, true)
Z
zhangkaihuo 已提交
184

185 186
- backward_api : sin_grad
  forward : sin(Tensor x) -> Tensor(out)
Z
zhangkaihuo 已提交
187
  args : (Tensor x, Tensor out_grad)
188 189
  output : Tensor(x_grad)
  kernel :
190 191 192 193 194 195 196 197 198 199
    func : sin_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sin_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : sinh_grad
  forward : sinh(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : sinh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sinh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
200

201 202 203 204 205 206 207
- backward_api : softmax_grad
  forward : softmax(Tensor x, int axis=-1) -> Tensor(out)
  args : (Tensor out, Tensor out_grad, int axis)
  output : Tensor(x_grad)
  kernel :
    func : softmax_csr_grad{sparse_csr, sparse_csr -> sparse_csr}

208 209
- backward_api : sqrt_grad
  forward : sqrt(Tensor x) -> Tensor(out)
210
  args : (Tensor out, Tensor out_grad)
211
  output : Tensor(x_grad)
212
  kernel :
213 214 215 216 217 218 219 220 221 222
    func : sqrt_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           sqrt_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

- backward_api : square_grad
  forward : square(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : square_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           square_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
223

224 225 226 227 228 229 230 231
- backward_api : subtract_grad
  forward : subtract(Tensor x, Tensor y) -> Tensor(out)
  args : (Tensor x, Tensor y, Tensor out_grad)
  output : Tensor(x_grad), Tensor(y_grad)
  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}

232 233 234 235 236 237 238 239
- backward_api : tan_grad
  forward : tan(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  kernel :
    func : tan_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           tan_csr_grad {sparse_csr, sparse_csr -> sparse_csr}

240 241
- backward_api : tanh_grad
  forward : tanh(Tensor x) -> Tensor(out)
242
  args : (Tensor out, Tensor out_grad)
243
  output : Tensor(x_grad)
244
  kernel :
245 246
    func : tanh_coo_grad {sparse_coo, sparse_coo -> sparse_coo},
           tanh_csr_grad {sparse_csr, sparse_csr -> sparse_csr}
247

248 249 250 251
- backward_api : values_grad
  forward : coo_values(Tensor x) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
252
  kernel :
253
    func : coo_values_grad{sparse_coo, dense-> sparse_coo}
254 255 256 257

- backward_api: fused_attention_grad
  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)
258
  output : Tensor(query_grad), Tensor(key_grad), Tensor(value_grad)
259 260
  kernel :
    func : fused_attention_csr_grad{dense, dense, dense, sparse_csr, dense -> dense, dense, dense}
261 262
    layout : softmax
    data_type: query