| 196 | } |
| 197 | |
| 198 | std::shared_ptr<Tensor> TensorImplDft::slice( |
| 199 | const std::vector<size_t>& start, const std::vector<size_t>& end, |
| 200 | const std::vector<size_t>& step) { |
| 201 | Layout layout; |
| 202 | mgb::TensorLayout layout_mge; |
| 203 | if (is_host()) { |
| 204 | layout_mge = m_host_tensor->layout(); |
| 205 | layout = to_lite_layout(m_host_tensor->layout()); |
| 206 | } else { |
| 207 | layout_mge = m_dev_tensor->layout(); |
| 208 | layout = to_lite_layout(m_dev_tensor->layout()); |
| 209 | } |
| 210 | |
| 211 | size_t length = start.size(); |
| 212 | LITE_ASSERT( |
| 213 | length == end.size() && length <= layout.ndim, |
| 214 | "The start and end must be the same size and less than layout " |
| 215 | "ndim."); |
| 216 | std::vector<mgb::Slice> slices; |
| 217 | if (step.size()) { |
| 218 | LITE_ASSERT(length == step.size(), "The start and step must be the same size."); |
| 219 | for (size_t i = 0; i < length; i++) { |
| 220 | slices.push_back(mgb::Slice{start[i], end[i], step[i]}); |
| 221 | } |
| 222 | } else { |
| 223 | for (size_t i = 0; i < length; i++) { |
| 224 | slices.push_back(mgb::Slice{start[i], end[i]}); |
| 225 | } |
| 226 | } |
| 227 | auto subspec = mgb::SubTensorSpec::make_from_offset_elem(layout_mge, 0); |
| 228 | size_t axis = 0; |
| 229 | for (auto&& i : slices) { |
| 230 | subspec.merge_with(i.apply(subspec.layout(), axis)); |
| 231 | axis++; |
| 232 | } |
| 233 | auto ret = std::make_shared<Tensor>(); |
| 234 | auto& impl = TensorHelper::implement(ret)->cast_final_safe<TensorImplDft>(); |
| 235 | if (is_host()) { |
| 236 | *impl.m_host_tensor = m_host_tensor->sub(subspec); |
| 237 | } else { |
| 238 | impl.m_dev_tensor = |
| 239 | std::make_shared<mgb::DeviceTensorND>(m_dev_tensor->sub(subspec)); |
| 240 | impl.m_host_tensor = nullptr; |
| 241 | } |
| 242 | LITE_ASSERT(is_host() == impl.is_host()); |
| 243 | return ret; |
| 244 | } |
| 245 | |
| 246 | void TensorImplDft::fill_zero() { |
| 247 | if (is_host()) { |