Helper for Split() that returns the dimension in the slice.
| 349 | |
| 350 | // Helper for Split() that returns the dimension in the slice. |
| 351 | static inline int GetDimensionInSlice(const int dim, const int split_size, |
| 352 | const int residual) { |
| 353 | DCHECK_GT(split_size, 0); |
| 354 | DCHECK_GE(dim, 0); |
| 355 | if (residual == 0) return dim % split_size; |
| 356 | const int offset = residual * (split_size + 1); |
| 357 | if (dim < offset) { |
| 358 | return dim % (split_size + 1); |
| 359 | } else { |
| 360 | return (dim - offset) % split_size; |
| 361 | } |
| 362 | } |
| 363 | |
| 364 | // Helper for Split() that returns the shape given a slice index. |
| 365 | static inline int GetSliceShape(const int slice_index, const int split_size, |