}

// Finds AssignVariableOps that can be moved into a parallel_execute region and
// moves them. Leftover parallel_execute results that were used by the
// 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;