message: Fixed tf.reduce_sum usage on 2-D tensors.
diff: @@ -138,12 +138,32 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) {
  memcpy(output_float_data.data(), (*input_float_data).data(),
  output_buffer_size * sizeof(output_float_data[0]));
  } else if (unary_op->type == OperatorType::kTensorFlowSum) {
  - // At the moment only full reduction across all dimensions is supported.
  + CHECK_EQ(unary_op->inputs.size(), 2) << "Sum needs 2 inputs";
  + if (!IsConstantParameterArray(*model, unary_op->inputs[1])) {
  + AddMessageF("Axis input is non-constant");
  + return false;
  + }
  + auto& axis_array = model->GetArray(unary_op->inputs[1]);
  + CHECK(axis_array.data_type == ArrayDataType::kInt32);
  + int axis = axis_array.GetBuffer<ArrayDataType::kInt32>().data[0];
  + CHECK_LT(axis, input_shape.dimensions_count()) << "Axis out of bounds";
  +
  + // We currently only handle reduction on axis 0.
  + CHECK_EQ(axis, 0) << "Only reduction along axis 0 is supported";
  + // We currently only handle 1-D and 2-D input tensors.
  + CHECK_LE(input_shape.dimensions_count(), 2) << "Rank >2 not yet supported";
  + // We only support keep_dims=true; shape prop will need to change otherwise.
  + auto sum_op = static_cast<const TensorFlowSumOperator*>(unary_op);
  + CHECK(sum_op->keep_dims) << "Only keep_dims=true is supported";
  +
  + std::vector<int> indices(input_shape.dimensions_count());
  + for (int i = 0; i < input_shape.dims(1); ++i) {
  + indices[1] = i;
  float sum = 0.f;
  - for (int i = 0; i < input_buffer_size; i++) {
  - sum += (*input_float_data)[i];
  + for (int j = 0; j < input_shape.dims(0); ++j) {
  + indices[0] = j;
  + sum += (*input_float_data)[Offset(input_shape, indices)];
  }
  - for (int i = 0; i < output_buffer_size; ++i) {
  output_float_data[i] = sum;
  }
  } else if (unary_op->type == OperatorType::kTensorFlowMin) {
  