reduce_op.getLoc(), reduce_op->getResult(0).getType(), operand,
          reduction_indices,
          /*keep_dim=*/rewriter.getBoolAttr(false));
      auto tf_argreduce_op = rewriter.create<ArgReduce>(
          reduce_op.getLoc(), reduce_op->getResult(1).getType(), operand,
          reduction_indices);

      rewriter.replaceOp(reduce_op, {tf_reduce_op, tf_argreduce_op});
    } else {

std::optional<bool> IsReduceOpLegal(mhlo::ReduceOp reduce_op) {
  if (succeeded(MatchReduceToArgMinMaxType1(reduce_op, true, true)) ||
      succeeded(MatchReduceToArgMinMaxType1(reduce_op, false, true)) ||
      succeeded(MatchReduceToArgMinMaxType1(reduce_op, true, false)) ||
      succeeded(MatchReduceToArgMinMaxType1(reduce_op, false, false)) ||
      succeeded(MatchReduceToArgMinMaxType2(reduce_op, false)) ||

        getInvertPermutedAttr(reduce_op.getWindowDimensions());
    DenseI64ArrayAttr new_window_strides =
        getInvertPermutedAttr(reduce_op.getWindowStrides());
    DenseI64ArrayAttr new_base_dilations =
        getInvertPermutedAttr(reduce_op.getBaseDilations());
    DenseI64ArrayAttr new_win_dilations =
        getInvertPermutedAttr(reduce_op.getWindowDilations());

      return failure();
    }

    StringRef reduce_op = all_reduce.getReduceOp();

    StringRef merge_op, final_op;
    if (reduce_op == "Add") {
      merge_op = "Add";
      final_op = "Id";
    } else if (reduce_op == "Mul") {
      merge_op = "Mul";
      final_op = "Id";
    } else if (reduce_op == "Min") {
      merge_op = "Min";
      final_op = "Id";

  // CHECK: "mhlo.all_reduce"
  // CHECK{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
  // CHECK-NOT: channel_handle
  // CHECK: mhlo.add
  %0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
  func.return %0 : tensor<f32>
}

// -----

// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 0

  if (!return_op) return failure();
  if (return_op.getNumOperands() != 1) return failure();

  ReductionOp reduce_op = dyn_cast_or_null<ReductionOp>(
      return_op.getOperands().front().getDefiningOp());
  if (!reduce_op) return failure();
  if (reduce_op.getLhs() != body.getArgument(0) ||
      reduce_op.getRhs() != body.getArgument(1))
    return failure();

  return success();
}

  // CHECK-NEXT:   "tf_device.cluster"() ({
  // CHECK-NEXT:     %[[GROUP:.*]] = "tf.Const"()
  // CHECK-SAME:       [0, 1, 2, 3]
  // CHECK-NEXT:     %[[REDUCED:.*]] = "tf.XlaAllReduce"(%[[ID]], %[[GROUP]]) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x4xi32>) -> tensor<f32>
  // CHECK-NEXT:     "tf.OpA"(%[[REDUCED]]) : (tensor<f32>) -> ()
  tf_device.replicate {n = 4 : i32} {
    "tf_device.cluster"() ({

  // This function tries to match that the "mhlo::ReduceOp" only has one
  // operand, one init_value and one result.
  LogicalResult MatchReduceOpOperand(mhlo::ReduceOp reduce_op) const {
    if (reduce_op.getInputs().size() != 1 ||
        reduce_op.getInitValues().size() != 1 ||
        reduce_op.getResults().size() != 1)
      return failure();

    if (!mlir::isa<RankedTensorType>(reduce_op.getInputs()[0].getType()))

func.func private @write_chain_func(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {
  %cst = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
  %0 = "tf.XlaAllReduce"(%arg0, %cst) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<1xf32>, tensor<1x2xi32>) -> tensor<1xf32>
  return %0 : tensor<1xf32>

          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),
                                        block_arg, group_assignment, reduce_op,
                                        mode);
}