diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
index 2ceda0d6360e..9eac12d4b56b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
@@ -317,20 +317,19 @@ static bool isIdentityMapWithZeros(AffineMap map) {
   if (map.isEmpty())
     return false;
   unsigned dimsSeen = 0;
-  for (auto result : map.getResults()) {
-    bool isValidExpr = TypeSwitch<AffineExpr, bool>(result)
-                           .Case<AffineDimExpr>([&dimsSeen](auto dimExpr) {
-                             if (dimExpr.getPosition() != dimsSeen)
-                               return false;
-                             dimsSeen++;
-                             return true;
-                           })
-                           .Case<AffineConstantExpr>([](auto constExpr) {
-                             return constExpr.getValue() == 0;
-                           })
-                           .Default([](AffineExpr) { return false; });
-    if (!isValidExpr)
+  for (AffineExpr result : map.getResults()) {
+    if (auto dimExpr = dyn_cast<AffineDimExpr>(result)) {
+      if (dimExpr.getPosition() != dimsSeen) {
+        return false;
+      }
+      dimsSeen++;
+    } else if (auto constExpr = dyn_cast<AffineConstantExpr>(result)) {
+      if (constExpr.getValue() != 0) {
+        return false;
+      }
+    } else {
       return false;
+    }
   }
   return dimsSeen == map.getNumDims();
 }