Add ConnectivityFields to handler

GridTools · Dec 11, 2024 · 9eb29be · 9eb29be
1 parent c9c5c8e
commit 9eb29be
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py
@@ -20,7 +20,7 @@
 from gt4py.eve import utils
 from gt4py.eve.utils import content_hash
 from gt4py.next import NeighborTableOffsetProvider, backend, common, config
-from gt4py.next.embedded.nd_array_field import NumPyArrayField
+from gt4py.next.embedded import nd_array_field
 from gt4py.next.iterator import ir as itir
 from gt4py.next.otf import arguments
 from gt4py.next.otf import recipes, stages, workflow
@@ -115,6 +115,8 @@ def handle_invalid_type(
 
 type_handlers_connectivity_args = {
     NeighborTableOffsetProvider: handle_connectivity,
+    nd_array_field.NumPyArrayConnectivityField: handle_connectivity,
+    nd_array_field.CuPyArrayConnectivityField:  handle_connectivity,
     common.Dimension: handle_dimension,
 }