// If ShapeOp operand is replicate tensor block argument, replace with the
  // associated first replica operand.
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(input)) {
    if (block_arg.getOwner() != replicate_block) return;

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

    return;
  }

  Operation* input_def = input.getDefiningOp();

      auto operand = SkipIdentityAndReadVariable(entry.value());
      auto block_arg = mlir::dyn_cast<BlockArgument>(operand);
      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {
        // Only mirrored args of ReplicateOp can be annotated.
        if (mirrored_replicate_args.count(block_arg.getArgNumber()) == 0) {
          continue;
        }

      return true;
    }
    if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(arg)) {
      // Skip the folding logic if the block argument is not from the function
      // arguments. This can happen when the argument is from a while loop.
      if (block_arg.getParentRegion() != &func_op.getRegion()) {
        return false;
      }

  }
  seen.insert(v);

  if (auto blockArg = mlir::dyn_cast<BlockArgument>(v)) {
    Operation *op = blockArg.getOwner()->getParentOp();

    // If we're in the first block, then the argument to that block is the
    // one we're looking for.
    if (auto func = cast<func::FuncOp>(op)) {
      if (blockArg.getOwner()->isEntryBlock()) {
        *out = blockArg;
        return success();
      }
    }

}

// Finds op that created a given value. If the value is a BlockArgument, this
// 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();
}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {

};

// Finds op that created a given value. If the value is a BlockArgument, this
// 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();
}

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {

};

// Finds op that created a given value. If the value is a BlockArgument, this
// 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.

// These Value(s) are taken to be the arguments of the new raised func.
// An operand dependency is a Value referenced anywhere in an Op
// that is defined above the Op. All SSA Values are assigned/defined in a
// BlockArg or as a result of an Operation.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
    const llvm::SetVector<Operation*>& partition_ops);

// Similar to `AccumulateOperandsDefinedAbove()`, computes the Value(s) that are