Gets offset of index if reducing on axis. When reducing, the flattened offset will not change, if the input index changes on the given axis. For example, if you have a 3D tensor and you are reducing to 2D by eliminating axis 0, then index (0, 1, 2) and index (1, 1, 2) will map to the same flattened offset. TODO(kanlig): uses Dims to represent dimensions.
| 349 | // offset. |
| 350 | // TODO(kanlig): uses Dims to represent dimensions. |
| 351 | inline size_t ReducedOutputOffset(const int num_dims, const int* dims, |
| 352 | const int* index, const int num_axis, |
| 353 | const int* axis) { |
| 354 | if (num_dims == 0) { |
| 355 | return 0; |
| 356 | } |
| 357 | TFLITE_DCHECK(dims != nullptr); |
| 358 | TFLITE_DCHECK(index != nullptr); |
| 359 | size_t offset = 0; |
| 360 | for (int idx = 0; idx < num_dims; ++idx) { |
| 361 | // if we need to skip this axis |
| 362 | bool is_axis = false; |
| 363 | if (axis != nullptr) { |
| 364 | for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) { |
| 365 | if (idx == axis[axis_idx]) { |
| 366 | is_axis = true; |
| 367 | break; |
| 368 | } |
| 369 | } |
| 370 | } |
| 371 | if (!is_axis) { |
| 372 | offset = offset * static_cast<size_t>(dims[idx]) + |
| 373 | static_cast<size_t>(index[idx]); |
| 374 | } |
| 375 | } |
| 376 | return offset; |
| 377 | } |
| 378 | |
| 379 | inline int Offset(const RuntimeShape& shape, int i0, int i1, int i2, int i3) { |
| 380 | TFLITE_DCHECK_EQ(shape.DimensionsCount(), 4); |