return true;
  };

  func_op.walk([&should_fold_op_func](TF::IdentityOp op) {
    StringAttr op_device = op->getAttrOfType<StringAttr>(kDeviceAttr);
    if (should_fold_op_func(op.getOperand(), op_device)) {
      op.replaceAllUsesWith(op.getOperand());
      op.erase();
    }
    return WalkResult::advance();
  });

  func_op.walk([&should_fold_op_func](TF::IdentityNOp op) {