| 75 | } |
| 76 | |
| 77 | Status ReductionHelper::Simplify(const Tensor& data, const Tensor& axis, |
| 78 | const bool keep_dims) { |
| 79 | // bitmap[i] indicates whether to reduce data along i-th axis. |
| 80 | gtl::InlinedVector<bool, 4> bitmap(data.dims(), false); |
| 81 | if (axis.dtype() == DT_INT32) { |
| 82 | TF_RETURN_IF_ERROR(SimplifyHelper<int32>(data, axis, bitmap)); |
| 83 | } else { |
| 84 | TF_RETURN_IF_ERROR(SimplifyHelper<int64>(data, axis, bitmap)); |
| 85 | } |
| 86 | // Output tensor's dim sizes. |
| 87 | out_shape_.clear(); |
| 88 | for (int i = 0; i < data.dims(); ++i) { |
| 89 | if (!bitmap[i]) { |
| 90 | // If we are not reducing along dimension i. |
| 91 | out_shape_.push_back(data.dim_size(i)); |
| 92 | } else if (keep_dims) { |
| 93 | // We are reducing along dimension i, but we want to keep the |
| 94 | // same number of dimensions, so we set the dimension of i to |
| 95 | // '1'. |
| 96 | out_shape_.push_back(1); |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | // Depending on bitmap[i] and bitmap[i-1], we can collapse axis of |
| 101 | // the input data before doing the reduction on the resulting |
| 102 | // tensor. The shape of the reduction is a reshape of the final |
| 103 | // output. |
| 104 | |
| 105 | // We'll skip the leading 1s. |
| 106 | int dim_index = 0; |
| 107 | for (; dim_index < data.dims(); ++dim_index) { |
| 108 | if (data.dim_size(dim_index) != 1) break; |
| 109 | } |
| 110 | if (dim_index >= data.dims()) { |
| 111 | // Special case. The input is essentially a scalar. |
| 112 | reduce_first_axis_ = true; |
| 113 | } else { |
| 114 | // Starting from the (dim_index)-th dimension, dimensions |
| 115 | // alternates between runs that need to be reduced and runs that |
| 116 | // don't. |
| 117 | // |
| 118 | // NOTE: If a dimension has size 1, we group it as the current |
| 119 | // run so that we can minimize the number of runs. |
| 120 | // |
| 121 | // E.g., when we want to reduce a tensor of shape [2, 1, 3, 1, |
| 122 | // 5] by axes = [1, 4], we should treat the tensor as a [6, 5] |
| 123 | // and reduce by axes = [1] (i.e., the output is shape [6]). |
| 124 | reduce_first_axis_ = bitmap[dim_index]; |
| 125 | data_reshape_.push_back(data.dim_size(dim_index)); |
| 126 | ++dim_index; |
| 127 | for (; dim_index < data.dims(); ++dim_index) { |
| 128 | const auto size = data.dim_size(dim_index); |
| 129 | if (size == 1) { |
| 130 | bitmap[dim_index] = bitmap[dim_index - 1]; |
| 131 | } |
| 132 | if (bitmap[dim_index - 1] != bitmap[dim_index]) { |
| 133 | // Starts a new run of reduce or !reduce. |
| 134 | data_reshape_.push_back(size); |