block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

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

    auto yield = builder->create<tf_executor::YieldOp>(
        terminator->getLoc(), terminator->getOperands());
    terminator->erase();

    // Create new island for each region.
    builder->setInsertionPoint(island_op);
    auto execute_island = builder->create<tf_executor::IslandOp>(
        island_op.getLoc(), yield.getOperandTypes(),
        island_op.getControl().getType(), island_op.getControlInputs());

    auto *new_ewise_op =
        rewriter.create(op->getLoc(), op->getName().getIdentifier(),
                        {tpose_arg1.getOperand(0), tpose_arg2.getOperand(0)},
                        new_out_type, op->getAttrs());

    // Apply original tranpose to output of ewise op.
    auto out_tpose_op = rewriter.create<TFL::TransposeOp>(
        new_ewise_op->getLoc(), op->getResult(0).getType(),

  builder.create<tf_executor::YieldOp>(terminator.getLoc(),
                                       terminator.getOperands());
  terminator.erase();

  builder.setInsertionPoint(island_op);
  IRMapping mapping;
  for (int i : llvm::seq<int>(0, num_replicas)) {
    // Create new island for replica.
    auto replica = builder.create<tf_executor::IslandOp>(
        island_op.getLoc(), output_types, control_type, replica_inputs);

  tf_executor::YieldOp yield = island.GetYield();
  if (yield.getNumOperands() == 0) {
    builder.create<TF::NoOp>(island.getLoc(), TypeRange{}, ValueRange{});
  } else if (yield.getNumOperands() == 1) {
    Value operand = yield.getOperand(0);
    auto identity = builder.create<TF::IdentityOp>(island.getLoc(),
                                                   operand.getType(), operand);
    yield.setOperand(0, identity.getOutput());

  SmallVector<Value, 4> tuple_operands;
  for (auto idx : llvm::seq<int32_t>(0, end))
    tuple_operands.push_back(
        builder.create<GetTupleElementOp>(value.getLoc(), value, idx)
            .getResult());

  return CreateTuple(builder, value.getLoc(), tuple_operands);
}

// Replaces uses of `values` with `replacements`. If `flatten_tuple` is false,

    FilterOutBlockArgControlDep(op->getOperands().drop_front(), new_operands);

    auto source_op = builder.create<tf_executor::NextIterationSourceOp>(
        op->getLoc(), op->getOperand(0).getType());
    builder.create<tf_executor::NextIterationSinkOp>(
        op->getLoc(), source_op.getToken(), /*input=*/op->getOperand(0),
        /*controlInputs=*/new_operands);
    op->replaceAllUsesWith(

std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {
  std::optional<QuantizationUnitLoc::QuantizationUnit> quant_unit =
      FindQuantizationUnitFromLoc(call_op->getLoc());
  return std::make_pair(quant_unit->func_name(), quant_unit->node_name());
}

std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::XlaCallModuleOp call_op, const FlatSymbolRefAttr &f_attr) {

        cluster_func->getLoc(), partitioned_output_types, result,
        partitioned_input.getPartitionDimsAttr(),
        partitioned_input.get_XlaShardingAttr());
    for (auto [i, value] : llvm::enumerate(partitioned_output.getOutput())) {
      const auto& resource = packed_input ? inputs[0] : inputs[i];
      builder.create<TF::AssignVariableOp>(
          assign_var->getLoc(), /*resource=*/resource, /*value=*/value);

        maxs.push_back(rmax);
      }
      quant_type = quantfork::fakeQuantAttrsToType(
          op.getLoc(), num_bits, *op.getAxis(), mins, maxs, narrow_range,
          expressed, is_signed);
      if (legacy_float_scale) {
        quant_type = DownCastScale(quant_type, mins, maxs, op->getLoc());
      }
    } else if (auto stats =
                   op.getLayerStats().dyn_cast<DenseFPElementsAttr>()) {