| 45 | } |
| 46 | |
| 47 | gtl::InlinedVector<int32, 8> ReductionHelper::permutation() { |
| 48 | const int dims = data_reshape_.size(); |
| 49 | const int unreduced_dims = (dims + !reduce_first_axis_) / 2; |
| 50 | gtl::InlinedVector<int32, 8> perm(dims); |
| 51 | for (int i = 0; i < unreduced_dims; i++) { |
| 52 | perm[i] = 2 * i + reduce_first_axis_; |
| 53 | } |
| 54 | for (int i = unreduced_dims; i < dims; i++) { |
| 55 | perm[i] = 2 * (i - unreduced_dims) + !reduce_first_axis_; |
| 56 | } |
| 57 | return perm; |
| 58 | } |
| 59 | |
| 60 | template <typename Tperm> |
| 61 | Status SimplifyHelper(const Tensor& data, const Tensor& axis, |