=============Matrix operations============================================
| 1186 | |
| 1187 | // =============Matrix operations============================================ |
| 1188 | Tensor Average(const Tensor &M, int axis) { |
| 1189 | // operator/ only has implementation for float scalar type, hence it is |
| 1190 | // necessary to cast the denominator to a float. |
| 1191 | // TODO(wangwei) implement function for cast scalar type involved in Tensor |
| 1192 | // functions. E.g., |
| 1193 | // template<S, D> |
| 1194 | // D CastTo(S x) { |
| 1195 | // return D(x); |
| 1196 | // } |
| 1197 | // for speical types, e.g., fp16: |
| 1198 | // tempalte<> |
| 1199 | // fp16 CastType(float x) { |
| 1200 | // .... |
| 1201 | // } |
| 1202 | if (axis == 0) { |
| 1203 | return Sum(M, 0) / (1.0f * M.shape(0)); |
| 1204 | } else if (axis == 1) { |
| 1205 | return Sum(M, 1) / (1.0f * M.shape(1)); |
| 1206 | } else { |
| 1207 | LOG(FATAL) << "Not currently support Sum over axis = " << axis; |
| 1208 | } |
| 1209 | } |
| 1210 | // TODO(wangwei) conside async exec |
| 1211 | template <> |
| 1212 | float Sum<float>(const Tensor &in) { |