message: [XLA] Correctly handle variadic reductions with a single output in the evaluator
diff: @@ -2276,11 +2276,11 @@ static bool IsScalarAdd(HloComputation* computation) {
  // (until the reduction is completed, the output element is also used as
  // an accumulator).
  static StatusOr<bool> PerformReductionStep(
  - absl::Span<const int64> input_index, absl::Span<const int64> output_index,
  + bool is_tuple, absl::Span<const int64> input_index,
  + absl::Span<const int64> output_index,
  absl::Span<const Literal* const> input_args, absl::Span<Literal> results,
  HloComputation* computation, HloEvaluator* embedded_evaluator) {
  int num_args = results.size();
  - bool is_tuple = num_args > 1;
  absl::InlinedVector<Literal, 1> arg_values;
  arg_values.reserve(num_args);
  @@ -2330,7 +2330,7 @@ static StatusOr<bool> PerformReductionStep(
  }
  static StatusOr<bool> GenerateReduceOutputElement(
  - absl::Span<const int64> output_index,
  + bool is_tuple, absl::Span<const int64> output_index,
  absl::Span<const Literal* const> init_values,
  absl::Span<const Literal* const> input_args, absl::Span<Literal> results,
  @@ -2340,7 +2340,6 @@ static StatusOr<bool> GenerateReduceOutputElement(
  absl::Span<const int64> arg_dim_steps,
  absl::Span<const int64> arg_dim_counts,
  absl::Span<const int64> result_to_arg_index) {
  - bool is_tuple = results.size() > 1;
  bool use_fast_add = ShapeUtil::ElementIsFloating(init_values[0]->shape()) &&
  IsScalarAdd(function) && !is_tuple;
  @@ -2375,8 +2374,9 @@ static StatusOr<bool> GenerateReduceOutputElement(
  TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus(
  arg_shape, base, arg_dim_counts, arg_dim_steps,
  [&](absl::Span<const int64> input_index) {
  - return PerformReductionStep(input_index, output_index, input_args,
  - results, function, embedded_evaluator);
  + return PerformReductionStep(is_tuple, input_index, output_index,
  + input_args, results, function,
  + embedded_evaluator);
  }));
  return true;
  }
  @@ -2453,9 +2453,9 @@ Status HloEvaluator::HandleReduce(HloInstruction* instr) {
  TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus(
  output_shape, [&](absl::Span<const int64> output_index) {
  return GenerateReduceOutputElement(
  - output_index, init_values, input_args, absl::Span<Literal>(results),
  - function, &embedded_evaluator, arg_dim_steps, arg_dim_counts,
  - result_to_arg_index);
  + is_tuple, output_index, init_values, input_args,
  + absl::Span<Literal>(results), function, &embedded_evaluator,
  + arg_dim_steps, arg_dim_counts, result_to_arg_index);
  }));
  if (is_tuple) {
  