Fix some errors in orchestration decorator

C2SM · Jan 13, 2025 · f0eca78 · f0eca78
1 parent d15c2f6
commit f0eca78
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
@@ -388,7 +388,7 @@ def test_verify_diffusion_init_against_savepoint(
         (dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"),
     ],
 )
-@pytest.mark.parametrize("ndyn_substeps, orchestration", [(2, False) , (2, True)])
+@pytest.mark.parametrize("ndyn_substeps, orchestration", [(2, False), (2, True)])
 def test_run_diffusion_single_step(
     savepoint_diffusion_init,
     savepoint_diffusion_exit,

diff --git a/model/common/src/icon4py/model/common/orchestration/decorator.py b/model/common/src/icon4py/model/common/orchestration/decorator.py
@@ -638,8 +638,8 @@ def _concretize_symbols_for_dace_structure(dace_cls, orig_cls):
 
         return {
             **{
-                "CellDim_sym": grid.offset_providers["C2E"].table.shape[0],
-                "EdgeDim_sym": grid.offset_providers["E2C"].table.shape[0],
+                "CellDim_sym": grid.offset_providers["C2E"].ndarray.shape[0],
+                "EdgeDim_sym": grid.offset_providers["E2C"].ndarray.shape[0],
                 "KDim_sym": grid.num_levels,
             },
             **concretize_symbols_for_dace_structure,

diff --git a/model/common/src/icon4py/model/common/orchestration/halo_exchange.py b/model/common/src/icon4py/model/common/orchestration/halo_exchange.py
@@ -282,7 +282,7 @@ def __sdfg__(self, *args, **kwargs) -> dace.SDFG:
         sdfg = dace.SDFG("DummyNestedSDFG")
         state = sdfg.add_state()
 
-        sdfg.add_scalar(name="__return", dtype=dace.int32)
+        sdfg.add_array(name="__return", shape=[1], dtype=dace.int32)
 
         tasklet = dace.sdfg.nodes.Tasklet(
             "DummyNestedSDFG",