for (const auto& idx_result : llvm::enumerate(island.getOutputs())) {
      Value result = idx_result.value();

      bool has_external_use = false;
      for (OpOperand& use : llvm::make_early_inc_range(result.getUses())) {
        if (wrapped_ops.count(use.getOwner()))
          use.set(yield_op.getOperand(idx_result.index()));
        else
          has_external_use = true;
      }
      if (has_external_use) {

llvm::SmallVector<Operation*> GetSuccessors(Operation* op) {
  llvm::SmallVector<Operation*> successors;
  for (auto result : op->getParentOp()->getOpResults()) {
    for (auto& use : result.getUses()) {
      auto succ = use.getOwner();
      successors.push_back(succ);
    }
  }
  return successors;
}
// Gets the predecessors of an op wrapped in tf_executor.island.

    // to it have the same shape, we need to follow use-def chain of ops that
    // (conceptually) modify it i.e., ops that take an input TensorList and
    // produce an output TensorList.
    for (auto& use : tensorlist.getUses()) {
      if (auto push = llvm::dyn_cast<TensorListPushBackOp>(use.getOwner())) {
        auto element_type =
            mlir::dyn_cast<RankedTensorType>(push.getTensor().getType());

      // with corresponding results of new parallel_execute op
      for (auto& use : llvm::make_early_inc_range(
               old_parallel_execute->getResult(output_index).getUses())) {
        use.set(new_parallel_execute->getResult(output_index));
      }
      continue;
    }

    int tpu_cluster_output_index = output_index - num_results_pre_cluster;
    const auto& output_sharding =

      has_dynamic_output = true;
      for (Operation* user : v.getUsers()) {
        if (!HasOutsideCompilationAncestor(user) &&
            !llvm::isa<mlir::tf_device::ReturnOp>(user))
          return true;
      }
    }
    if (!tensorflow::TypeValidForXLA(v.getType()))
      for (const Operation* user : v.getUsers())
        if (!llvm::isa<mlir::tf_device::ReturnOp>(user))

    llvm::SmallVector<Value, 4> next_values_to_visit;
    for (Value value_to_visit : values_to_visit) {
      if (!visited_values.insert(value_to_visit).second) continue;

      for (auto& use : value_to_visit.getUses()) {
        Operation* owner = use.getOwner();
        if (auto sharding = llvm::dyn_cast<mlir::TF::XlaShardingOp>(owner))
          return sharding.get_XlaSharding();

        // Only `mlir::func::CallOp` is supported as this requires knowing how
        // to rewrite arguments and results to a function.
        if (!isa<mlir::func::CallOp>(use.getUser())) continue;
        auto caller_parent_func =
            use.getUser()->getParentOfType<func::FuncOp>();
        if (!caller_parent_func) continue;

        FuncToRewrite func_to_rewrite{/*original=*/caller_parent_func,

    ResourceInfo info;
    info.data_type = GetResourceSubtype(resource);
    llvm::BitVector written_regions(op_->getNumRegions());
    bool unsupported_use = false;
    for (OpOperand& use : resource.getUses()) {
      Operation* user = use.getOwner();
      // If the user is not in one of the regions, we are not interested in it.
      // Since all the sub-regions within this region (i.e., regions attached to

  for (auto& use : uses_optional.value()) {
    if (llvm::isa<TF::WhileOp>(use.getUser())) {
      if (while_op != nullptr) {
        use.getUser()->emitOpError() << "multiple users of function.";
        return LogicalResult::failure();
      } else {
        while_op = llvm::cast<TF::WhileOp>(use.getUser());
      }
    } else {
      use.getUser()->emitOpError() << "non while use of function.";
      return LogicalResult::failure();

        !(qtype.isa<mlir::quant::UniformQuantizedType>() ||
          qtype.isa<mlir::quant::UniformQuantizedPerAxisType>())) {
      return;
    }
    for (auto& use : value.getUses()) {
      Operation* user = use.getOwner();
      if (user->hasTrait<mlir::OpTrait::IsTerminator>()) continue;

      auto affine_user = llvm::dyn_cast<mlir::AffineQuantizedOpInterface>(user);
      if (affine_user &&