// Move AssignVariableOp and update the value to be written to the
      // resource variable to be the non forwarded value from within the
      // parallel_execute region.
      assign_var.getOperation()->moveBefore(terminator);
      assign_var.getValueMutable().assign(
          terminator->getOperand(result.index()));
      results_to_remove.push_back(result.index());
    }

    rewrite |= !results_to_remove.empty();

        num_partitions_after += op_partition && (scheduled_something == false);
        scheduled_something = true;
        // Schedule the operation by moving it to the start.
        unscheduled_op.moveBefore(block, next_unscheduled_op);
        if (&unscheduled_op == &*next_unscheduled_op) {
          ++next_unscheduled_op;
        }
      }
    }

    // defined by tf_device.replicate.
    if (DescendantUsesVirtualDevice(virtual_device_list, &inner_op)) continue;

    if (IsOpReplicateInvariant(replicate_region, &inner_op))
      inner_op.moveBefore(replicate_op);
  }
}

void ReplicateInvariantOpHoistingPass::runOnOperation() {
  getOperation().walk(
      [](tf_device::ReplicateOp op) { HoistReplicateInvariantOps(op); });
}

  // Move all operations in cluster to newly created region, stripping their
  // "device" attribute since launch op already carries device information.
  Block* block = &region.front();
  for (Operation* op : c.ops) {
    op->moveBefore(block, block->end());
    op->removeAttr(builder->getStringAttr("device"));
  }

  // Get all escaped live-out values of region, they are used later to determine
  // return values and types of launch op.

  Block* block = &island.getBody().back();
  OpBuilder island_builder(original_island);
  island_builder.setInsertionPointToEnd(block);
  sub_op.replaceAllUsesWith(island.getOutputs());
  sub_op.moveBefore(block, block->begin());
  island_builder.create<tf_executor::YieldOp>(loc, sub_op.getResults());
  return island;
}

// Converts a single island into multiple islands (one for each op).

                            tpu_copy_with_dynamic_shape_op->getLoc(),
                            tpu_copy_with_dynamic_shape_op->getResults())
                        .getOperation();
  tpu_copy_with_dynamic_shape_op->moveBefore(return_op);
  return success();
}

// Update all the usage of tf_device.return op with launch op result.
void UpdateReturnOpResultWithLaunchOpResult(tf_device::LaunchOp* launch_op) {