Support new find_distributed_partition (#298)

* Support new find_distributed_partition inducer/pytato#393 changes the function signature. * pylint * flake8 * only catch one type of TypeError Co-authored-by: Matt Smith <[email protected]> * flake8 --------- Co-authored-by: Matt Smith <[email protected]>
inducer · May 9, 2023 · 3dc79c2 · 3dc79c2
1 parent 7380319
commit 3dc79c2
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/grudge/array_context.py b/grudge/array_context.py
@@ -232,7 +232,21 @@ def _dag_to_compiled_func(self, dict_of_named_arrays,
         self.actx._compile_trace_callback(self.f, "pre_find_distributed_partition",
                 dict_of_named_arrays)
 
-        distributed_partition = pt.find_distributed_partition(dict_of_named_arrays)
+        # https://github.com/inducer/pytato/pull/393 changes the function signature
+        try:
+            # pylint: disable=too-many-function-args
+            distributed_partition = pt.find_distributed_partition(
+                # pylint-ignore-reason:
+                # '_BasePytatoArrayContext' has no
+                # 'mpi_communicator' member
+                # pylint: disable=no-member
+                self.actx.mpi_communicator, dict_of_named_arrays)
+        except TypeError as e:
+            if "find_distributed_partition() takes 1 positional" in str(e):
+                distributed_partition = pt.find_distributed_partition(
+                    dict_of_named_arrays)
+            else:
+                raise
 
         if __debug__:
             # pylint-ignore-reason: