| 218 | } |
| 219 | |
| 220 | std::tuple<int, int, int> GetDimIndices(const DataLayout& layout, |
| 221 | const int data_dims) { |
| 222 | int depth_idx, batch_idx, spatial_idx; |
| 223 | switch (layout) { |
| 224 | case DataLayout::kYXBatchDepth: |
| 225 | depth_idx = data_dims - 1; |
| 226 | batch_idx = data_dims - 2; |
| 227 | spatial_idx = 0; |
| 228 | break; |
| 229 | |
| 230 | case DataLayout::kYXDepthBatch: |
| 231 | depth_idx = data_dims - 2; |
| 232 | batch_idx = data_dims - 1; |
| 233 | spatial_idx = 0; |
| 234 | break; |
| 235 | |
| 236 | case DataLayout::kBatchYXDepth: |
| 237 | depth_idx = data_dims - 1; |
| 238 | batch_idx = 0; |
| 239 | spatial_idx = 1; |
| 240 | break; |
| 241 | |
| 242 | case DataLayout::kBatchDepthYX: |
| 243 | case DataLayout::kBatchDepthYX4: |
| 244 | depth_idx = 1; |
| 245 | batch_idx = 0; |
| 246 | spatial_idx = 2; |
| 247 | break; |
| 248 | |
| 249 | default: |
| 250 | LOG(FATAL) << "Unknown layout " << layout; |
| 251 | } |
| 252 | |
| 253 | return std::make_tuple(depth_idx, batch_idx, spatial_idx); |
| 254 | } |
| 255 | |
| 256 | std::vector<int64> ReorderDims(const std::vector<int64>& input, |
| 257 | const DataLayout& from, const DataLayout& to) { |