sparse_api.yaml 4.5 KB
Newer Older
1 2 3 4 5 6 7 8 9
- api : add
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : add_coo_coo{sparse_coo -> sparse_coo},
           add_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : add_grad

10
- api : conv3d
Z
zhangkaihuo 已提交
11
  args : (Tensor x, Tensor kernel, int[] paddings, int[] dilations, int[] strides, int groups, bool subm)
12
  output : Tensor(out), Tensor(rulebook)
13
  kernel :
14
    func : sparse_conv3d{sparse_coo, dense -> sparse_coo, dense}
15
    layout : x
16
  intermediate : rulebook
17
  backward : conv3d_grad
18

19 20
- api : coo_to_dense
  args : (Tensor x)
21
  output : Tensor(out)
22 23 24
  invoke : to_dense_impl(x)
  backward : coo_to_dense_grad

25 26
- api : create_sparse_coo_tensor
  args : (Tensor values, Tensor indices, IntArray dense_shape)
27
  output : Tensor(out)
28
  kernel :
29
    func : sparse_coo_tensor{dense, dense -> sparse_coo}
30
    layout : values
31
    data_type : values
32 33
  backward : create_sparse_coo_tensor_grad

34 35 36 37 38
- api : dense_to_coo
  args : (Tensor x, int64_t sparse_dim)
  output : Tensor(out)
  invoke : to_sparse_coo_impl(x, sparse_dim)
  backward : dense_to_coo_grad
39

40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
- api : divide
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : divide_coo_coo{sparse_coo -> sparse_coo},
           divide_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : divide_grad

- api : multiply
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : multiply_coo_coo{sparse_coo -> sparse_coo},
           multiply_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : multiply_grad

58
- api : relu
59
  args : (Tensor x)
60
  output : Tensor(out)
61
  kernel :
62 63
    func : sparse_coo_relu{sparse_coo -> sparse_coo},
           sparse_csr_relu{sparse_csr -> sparse_csr}
64
    layout : x
65
  backward : relu_grad
66

67
- api : sin
68
  args : (Tensor x)
69
  output : Tensor(out@SparseCooTensor)
70
  kernel :
71 72
    func : sparse_coo_sin {sparse_coo -> sparse_coo},
           sparse_csr_sin {sparse_csr -> sparse_csr}
73
    layout : x
74
  backward : sin_grad
75

76 77 78 79 80 81 82 83
- api : softmax
  args : (Tensor x, int axis=-1)
  output : Tensor(out)
  kernel :
    func : softmax_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : softmax_grad

84
- api : sqrt
85
  args : (Tensor x)
86
  output : Tensor(out)
87
  kernel :
88 89
    func : sparse_coo_sqrt{sparse_coo -> sparse_coo},
           sparse_csr_sqrt{sparse_csr -> sparse_csr}
90
    layout : x
91
  backward : sqrt_grad
92

93 94 95 96 97 98 99 100 101
- api : subtract
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : subtract_coo_coo{sparse_coo -> sparse_coo},
           subtract_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : subtract_grad

102
- api : tanh
103
  args : (Tensor x)
104
  output : Tensor(out)
105
  kernel :
106 107
    func : sparse_coo_tanh{sparse_coo -> sparse_coo},
           sparse_csr_tanh{sparse_csr -> sparse_csr}
108
    layout : x
109
  backward : tanh_grad
110

111
- api : to_dense
112
  args : (Tensor x)
113
  output : Tensor(out)
114
  invoke : to_dense_impl(x)
115

116
- api : to_sparse_coo
117
  args : (Tensor x, int64_t sparse_dim)
118
  output : Tensor(out)
119
  invoke : to_sparse_coo_impl(x, sparse_dim)
120

121
- api : to_sparse_csr
122
  args : (Tensor x)
123
  output : Tensor(out)
124
  invoke : to_sparse_csr_impl(x)
Z
zhangkaihuo 已提交
125

126 127 128 129 130 131 132 133 134
- api : values
  args : (Tensor x)
  output : Tensor(out)
  kernel :
    func : coo_values{sparse_coo -> dense},
           csr_values{sparse_csr -> dense}
    layout : x
  backward : values_grad

135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
- api: full_like
  args : (Tensor x, Scalar value, DataType dtype=DataType::UNDEFINED)
  output : Tensor(out)
  kernel :
    func : coo_full_like{sparse_coo -> sparse_coo},
           csr_full_like{sparse_csr -> sparse_csr}
    layout : x
    data_type : dtype

- api: masked_matmul
  args : (Tensor x, Tensor y, Tensor mask)
  output : Tensor(out)
  kernel :
    func : csr_masked_matmul{dense, dense, sparse_csr -> sparse_csr}
    layout : x
  backward: masked_matmul_grad

- api: matmul
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : csr_dense_matmul{sparse_csr, dense -> dense},
           csr_csr_matmul{sparse_csr, sparse_csr -> sparse_csr},
           coo_dense_matmul{sparse_coo, dense -> dense},
           coo_coo_matmul{sparse_coo, sparse_coo -> sparse_coo}
    layout : x
  backward: matmul_grad

Z
zhangkaihuo 已提交
163 164
- api: maxpool
  args : (Tensor x, int[] kernel_sizes, int[] paddings, int[] dilations, int[] strides)
165
  output : Tensor(out), Tensor(rulebook)
Z
zhangkaihuo 已提交
166
  kernel :
167
    func : sparse_maxpool{sparse_coo -> sparse_coo, dense}
Z
zhangkaihuo 已提交
168 169 170
    layout : x
  intermediate : rulebook
  backward : sparse_maxpool_grad
171 172 173 174 175 176 177 178 179

- api: mv
  args : (Tensor x, Tensor vec)
  output : Tensor(out)
  kernel :
    func : mv_coo{sparse_coo, dense -> dense},
           mv_csr{sparse_csr, dense -> dense}
    layout : x
  backward: mv_grad