operand->get().getType() ==
              ShuffleRankedTensorType(ranked_tranpose_type, permutation)) {
        TransposeOp transpose = tranpose_op;
        transpose.getOperation()->moveBefore(op);
        transpose.setOperand(0, operand->get());
        transpose.setOperand(1, permutation_op);
        transpose_ops->erase(it);
        return transpose;
      }
    }
  }
  return nullptr;
}

  }
  Block* old_block = preprocess_op->getBlock();
  Block* launch_block = new Block;
  for (Operation& op : llvm::make_early_inc_range(old_block->getOperations())) {
    if (ops_to_move.contains(&op)) {
      op.moveBefore(launch_block, launch_block->end());
    }
  }

  // Create the new launch op after TPUCompileOp with the preprocessing ops.
  tf_device::LaunchOp moved_launch_op =

      if (read->isBeforeInBlock(peer)) continue;
      if (llvm::isa<ReadVariableOp>(peer)) continue;
      has_conflicting_write = true;
    }
    if (has_conflicting_write) continue;
    read->moveBefore(replicate);
  }
}

// Hoist `ReadVariableOp`s above the `tf_device.replicate`s.
void HoistBroadcastRead::runOnOperation() {
  FuncOp func = getOperation();

    // 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); });
}

                                                          cluster.getResults());
  }

  for (auto pi : partitioned_inputs) pi->moveBefore(result_op);

  cluster.getOperation()->moveBefore(result_op);

  return success();
}

// Forms clusters with ops of the same `_replication_info` attribute under a
// block.
//

                            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) {

  Block* launch_block = new Block;
  for (Operation* head_outside_compiled_op : head_outside_compiled_ops) {
    head_outside_compiled_op->removeAttr(kXlaOutsideCompilationAttr);
    head_outside_compiled_op->moveBefore(launch_block, launch_block->end());
  }

  mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
      builder, cluster, /*before=*/true, launch_block, host_device);

      // and packed args.
      arg.replaceAllUsesWith(
          new_replicate.GetBody().getArgument(arg.getArgNumber() + 1));
    }
  }
  for (auto& op : llvm::make_early_inc_range(replicate.GetBody())) {
    op.moveBefore(&new_replicate.GetBody(), new_replicate.GetBody().end());
  }
  replicate.replaceAllUsesWith(new_replicate);
  replicate.erase();
  return new_replicate;
}

          HasControlDependencyWithUnscheduledOp(op, block, unscheduled_ops)) {
        continue;
      }
      unscheduled_ops.erase(&op);
      if (Block::iterator(op) != last_scheduled_op)
        op.moveBefore(block, last_scheduled_op);
      else
        ++last_scheduled_op;
      scheduled_at_least_once = true;
    }
    if (!scheduled_at_least_once) return failure();
  }
  return success();
}

                                      Operation* after_op_r) {
  for (Operation& op : llvm::make_early_inc_range(device_cluster.GetBody())) {
    if (&op != device_cluster.GetBody().getTerminator()) {
      op.moveBefore(after_op_r);
    }
  }
}

// Move ops in the tmp host launch op to new host launch op
void MoveTmpLaunchOpToNewLaunchOp(mlir::tf_device::LaunchOp tmp_host_launch_op,