fix test

GridTools · Dec 5, 2024 · 38fd60c · 38fd60c
1 parent 7c115af
commit 38fd60c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
@@ -1707,7 +1707,9 @@ def impl(*iters: ItIterator):
 
 
 def _dimension_to_tag(domain: Domain) -> dict[Tag, range]:
-    return {k.value if isinstance(k, common.Dimension) else k: v for k, v in domain.items()}
+    return type(domain)(  # persist kind of domain
+        **{k.value if isinstance(k, common.Dimension) else k: v for k, v in domain.items()}
+    )
 
 
 def _validate_domain(domain: Domain, offset_provider_type: common.OffsetProviderType) -> None:

diff --git a/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py b/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py
@@ -27,21 +27,20 @@ def foo(inp):
     dtype=None,
 )
 
+I = gtx.Dimension("I")
+
 
 def test_deduce_domain():
     assert isinstance(_deduce_domain({}, {}), CartesianDomain)
     assert isinstance(_deduce_domain(UnstructuredDomain(), {}), UnstructuredDomain)
     assert isinstance(_deduce_domain({}, {"foo": connectivity}), UnstructuredDomain)
     assert isinstance(
-        _deduce_domain(CartesianDomain([("I", range(1))]), {"foo": connectivity}), CartesianDomain
+        _deduce_domain(CartesianDomain([(I, range(1))]), {"foo": connectivity}), CartesianDomain
     )
 
 
-I = gtx.Dimension("I")
-
-
 def test_embedded_error_on_wrong_domain():
-    dom = CartesianDomain([("I", range(1))])
+    dom = CartesianDomain([(I, range(1))])
 
     out = gtx.as_field([I], np.zeros(1))
     with pytest.raises(RuntimeError, match="expected 'UnstructuredDomain'"):