| 146 | } |
| 147 | |
| 148 | void Compute(OpKernelContext* ctx) override { |
| 149 | const Tensor& data = ctx->input(0); |
| 150 | const Tensor& axes = ctx->input(1); |
| 151 | VLOG(1) << "data shape: " << data.shape().DebugString(); |
| 152 | VLOG(1) << "axes : " << axes.SummarizeValue(10); |
| 153 | |
| 154 | ReductionHelper helper; |
| 155 | OP_REQUIRES_OK(ctx, helper.Simplify(data, axes, keep_dims_)); |
| 156 | CHECK_GE(helper.ndims(), 0); |
| 157 | |
| 158 | if (helper.ndims() == 0 || |
| 159 | (helper.ndims() == 1 && !helper.reduce_first_axis())) { |
| 160 | // Special case. Reduces nothing. It is unclear why this is |
| 161 | // necessary, but tests fail without it. Look into why this |
| 162 | // case occurs. |
| 163 | Tensor out; |
| 164 | if (!out.CopyFrom(data, helper.out_shape())) { |
| 165 | ctx->SetStatus(errors::Internal("Error during reduction copy.")); |
| 166 | } |
| 167 | ctx->set_output(0, out); |
| 168 | return; |
| 169 | } |
| 170 | |
| 171 | // We must allocate temp tensors using the same alloc attr as |
| 172 | // output(0) because it is returned as output(0) in the end. |
| 173 | const AllocatorAttributes alloc_attr = ctx->output_alloc_attr(0); |
| 174 | |
| 175 | // A temporary tensor whose size matches the size of the reduced |
| 176 | // output. |
| 177 | Tensor tmp_out; |
| 178 | OP_REQUIRES_OK( |
| 179 | ctx, ctx->allocate_temp(ctx->expected_output_dtype(0), |
| 180 | helper.out_reshape(), &tmp_out, alloc_attr)); |
| 181 | |
| 182 | typedef functor::ReduceFunctor<Device, Reducer> Functor; |
| 183 | Constants<Device> constants; |
| 184 | const Device& d = ctx->eigen_device<Device>(); |
| 185 | Reducer reducer; |
| 186 | |
| 187 | if (tmp_out.NumElements() == 0) { |
| 188 | // Nothing to do, fall through to final reshaping. |
| 189 | } else if (data.NumElements() == 0) { |
| 190 | // Degenerate reduction where the input is empty but the output is |
| 191 | // nonempty (thus tmp_out.NumElements() > 0), and we must fill the output |
| 192 | // with identity elements. Example: tf.reduce_sum(tf.zeros((0, 3)), [0]). |
| 193 | // Eigen sometimes crashes in this case, so we do it manually. |
| 194 | Functor::FillIdentity(d, tmp_out.flat<T>(), reducer); |
| 195 | } else if ((helper.ndims() == 1) && helper.reduce_first_axis()) { |
| 196 | // Reduce to a scalar. |
| 197 | Functor::Reduce(ctx, helper.out<T, 0>(&tmp_out), helper.in<T, 1>(data), |
| 198 | constants.kZero, reducer); |
| 199 | } else if ((helper.ndims() == 2) && helper.reduce_first_axis()) { |
| 200 | // Can be viewed as a reduction of a matrix along 1st dimension. |
| 201 | Functor::Reduce(ctx, helper.out<T, 1>(&tmp_out), helper.in<T, 2>(data), |
| 202 | constants.kZero, reducer); |
| 203 | } else if ((helper.ndims() == 2) && !helper.reduce_first_axis()) { |
| 204 | // Can be viewed as a reduction of a matrix along 2nd dimension. |
| 205 | Functor::Reduce(ctx, helper.out<T, 1>(&tmp_out), helper.in<T, 2>(data), |
nothing calls this directly
no test coverage detected