shape_op.setOperand(replicate_op.GetReplicaOperandForBlockArgument(
        block_arg, /*replica=*/0));

    return;
  }

  Operation* input_def = input.getDefiningOp();

  // If ShapeOp operand is a ReadVariableOp result where the ReadVariableOp
  // operand is a replicate resource block argument, replace ShapeOp with

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// Check if the TPUCopyWithDynamicShapeOp is valid.
// 1. The op should be wrapped inside a launch op.
// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {

  if (!converter.isLegal(convert_op.getOperand().getType())) {
    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
        convert_op.getOperand().getDefiningOp());
    return other_convert_op &&
           converter.isLegal(other_convert_op.getOperand().getType());
  } else if (!converter.isLegal(convert_op.getResult().getType())) {
    if (!convert_op.getResult().hasOneUse()) {

  for (int i = 0; i != num_return_operands; ++i) {
    auto returned_value = terminator->getOperand(i);
    Type returned_type = returned_value.getType();
    Operation* returned_op = returned_value.getDefiningOp();
    if (returned_op && llvm::isa<DequantizeOp>(returned_op)) {
      auto dequantize_op = llvm::cast<DequantizeOp>(returned_op);
      auto dequantize_input = dequantize_op.getInput();

Operation* GetInputOpWithOneUse(Operation* op, int opr_num) {
  if (opr_num >= op->getNumOperands()) return nullptr;
  auto opr = op->getOperand(opr_num);
  if (llvm::isa<BlockArgument>(opr)) return nullptr;
  auto* res = opr.getDefiningOp();
  if (!res->hasOneUse()) return nullptr;
  return res;
}

// Checks if the given operand of given operation refers to a splat constant
// with given val.
bool HasSplatArg(Operation* op, float val, int opr_num) {