}

void HandleFuncOp(Operation* op) {
  auto func = llvm::cast<func::FuncOp>(op);
  UpdateFuncType(func);
}

// Handles cast op between the first convolution and the block argument.
LogicalResult HandleCast(TF::CastOp cast_op, ArrayRef<int64_t> new_shape) {
  auto cast_input = cast_op.getX();
  // Update input type.
  auto transform_result_type =
      RankedTensorType::get(new_shape, getElementTypeOrSelf(cast_input));