| 1227 | } |
| 1228 | |
| 1229 | Tensor Sum(const Tensor &M, int axis) { |
| 1230 | if (axis == 0) { |
| 1231 | Tensor out(Shape{M.shape(1)}, M.device(), M.data_type()); |
| 1232 | SumRows(M, &out); |
| 1233 | return out; |
| 1234 | } else { |
| 1235 | CHECK_EQ(axis, 1) << "Not support Sum over axis = " << axis; |
| 1236 | Tensor out = Tensor(Shape{M.shape(0)}, M.device(), M.data_type()); |
| 1237 | SumColumns(M, &out); |
| 1238 | return out; |
| 1239 | } |
| 1240 | } |
| 1241 | |
| 1242 | Tensor SumAll(const Tensor &in) { |
| 1243 | Tensor out({(size_t)1}, in.device(), in.data_type()); |