<< "Cannot find function " << sym.getRootReference();
    }
    callees.push_back(callee);
  }
  return success();
}

bool HasSingleBlock(func::FuncOp func) {
  return func->getNumRegions() == 1 && func.getBody().hasOneBlock();
}

// This is a lengthy bit of code, since it has to recreate the operation.
// TODO(kramm): Move this under utils/ somewhere.
void EraseResults(Operation* op, llvm::BitVector erase) {
  assert(!op->getNumRegions());
  std::vector<Type> result_types;
  for (auto result : op->getResults()) {
    if (!erase[result.getResultNumber()]) {
      result_types.push_back(result.getType());
    }
  }
  OpBuilder builder(op);

// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
// TODO(b/284222084): Move executor dialect utilities to a new library.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
  if (func_op->getNumRegions() == 0 || func_op.getBody().empty()) return {};

  auto graph_op_range = func_op.front().without_terminator();
  if (llvm::hasSingleElement(graph_op_range)) {

// such AssignVariableOp are also pruned.
void SinkResourceWritesIntoParallelExecute(
    tf_device::ParallelExecuteOp parallel_execute) {
  bool rewrite = false;
  const int num_regions = parallel_execute.getNumRegions();
  llvm::SmallVector<Value, 4> results_to_remap;

  // Go through each region and find AssignVariableOps that can be moved into
  // the parallel_execute region. Result indices by region index are collected,

  // could be folded should have a custom folder instead of relying on the
  // TensorFlow folding hook.
  if (inst == nullptr || inst->getNumResults() == 0 ||
      inst->hasTrait<OpTrait::TF::NoConstantFold>() ||
      inst->getNumRegions() != 0 || !isMemoryEffectFree(inst)) {
    return false;
  }

  // If any of the result types are variants, don't try to constant fold them.
  // This creates opaque variant constants which lose information and would