// within other sub-regions attached to ops in this region. So the check
      // for whether a user is in one of the regions attached to this op is
      // straightforward.
      if (user->getParentRegion()->getParentOp() != op_) continue;

      // For functions, if the resource is used as a return operand, use that
      // as its result index.
      if (is_func && isa<func::ReturnOp>(user)) {

  Region* region = (*forward_pass_ops.begin())->getParentRegion();
  StringAttr replication_attr = GetReplicationAttr(*forward_pass_ops.begin());
  llvm::SmallVector<Operation*> checkset(forward_pass_ops.getArrayRef());
  checkset.append(backward_pass_ops.begin(), backward_pass_ops.end());
  for (Operation* op : checkset) {
    if (op->getParentRegion() != region) {

  // `mhlo.while` cond terminator does not need to be rewritten as it always
  // returns a tensor<i1> predicate value.
  if (auto while_parent = dyn_cast_or_null<WhileOp>(terminator->getParentOp()))
    if (terminator->getParentRegion() == &while_parent.getCond()) return;

  builder.setInsertionPoint(terminator);
  llvm::SmallDenseMap<Value, Value> rewritten_operands;
  auto new_results =

    RegionBranchPoint point) {
  if (auto whileOp =
          llvm::dyn_cast<WhileRegionOp>(this->getOperation()->getParentOp())) {
    if (&whileOp.getCond() == this->getOperation()->getParentRegion()) {
      // cut off the boolean (the condition itself) at the start
      return MutableOperandRange(
          this->getOperation(), 1,
          this->getOperation()->getOperands().size() - 1);
    }