use explicit isinstance() check instead of try/except around to_index_lambda()

Author: majosm

pytato/analysis/__init__.py

@@ -41,17 +41,19 @@
     Array,
     ArrayOrScalar,
     Concatenate,
+    DataWrapper,
     DictOfNamedArrays,
     Einsum,
     IndexBase,
     IndexLambda,
     IndexRemappingBase,
     InputArgumentBase,
     NamedArray,
+    Placeholder,
     ShapeType,
     Stack,
 )
-from pytato.diagnostic import CannotBeLoweredToIndexLambda
+from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
 from pytato.function import Call, FunctionDefinition, NamedCallResult
 from pytato.scalar_expr import (
     FlopCounter as ScalarFlopCounter,
@@ -75,7 +77,6 @@
 
     import pytools.tag
 
-    from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
     from pytato.loopy import LoopyCall
 
 __doc__ = """
@@ -776,10 +777,16 @@ def combine(self, *args: int) -> int:
         return sum(args)
 
     def _get_own_flop_count(self, expr: Array) -> int:
-        try:
-            nflops = self.scalar_flop_counter(to_index_lambda(expr).expr)
-        except CannotBeLoweredToIndexLambda:
-            nflops = 0
+        if isinstance(
+                expr,
+                (
+                    DataWrapper,
+                    Placeholder,
+                    NamedArray,
+                    DistributedRecv,
+                    DistributedSendRefHolder)):
+            return 0
+        nflops = self.scalar_flop_counter(to_index_lambda(expr).expr)
         if not isinstance(nflops, int):
             from pytato.scalar_expr import InputGatherer as ScalarInputGatherer
             var_names: set[str] = set(ScalarInputGatherer()(nflops))