Repository: GridTools/gt4py
Pull request #1317 (2023-08-11 12:41:31)

bug[next]: Fix ITIR nested tuple return tracing
While returning a tuple from an ITIR stencil could be successfully traced
```python
@fundef
def tuple_output1(inp1, inp2):
    return (deref(inp1), deref(inp2))
```
nested tuples were broken, which is fixed by this PR.
```python
@fundef
def tuple_of_tuple_output1(inp1, inp2, inp3, inp4):
    return (deref(inp1), deref(inp2)), (deref(inp3), deref(inp4))
```


Patch:
diff --git a/src/gt4py/next/iterator/tracing.py b/src/gt4py/next/iterator/tracing.py
index 6d4518f1..63b8c582 100644
--- a/src/gt4py/next/iterator/tracing.py
+++ b/src/gt4py/next/iterator/tracing.py
@@ -85,13 +85,13 @@ def _patch_Expr():
 
     @monkeypatch_method(Expr)
     def __call__(self, *args):
-        return FunCall(fun=self, args=[*make_node(args)])
+        return FunCall(fun=self, args=[make_node(arg) for arg in args])
 
 
 def _patch_FunctionDefinition():
     @monkeypatch_method(FunctionDefinition)
     def __call__(self, *args):
-        return FunCall(fun=SymRef(id=str(self.id)), args=[*make_node(args)])
+        return FunCall(fun=SymRef(id=str(self.id)), args=[make_node(arg) for arg in args])
 
 
 _patch_Expr()
@@ -114,7 +114,7 @@ def _f(fun, *args):
         fun = _s(fun)
 
     args = [trace_function_argument(arg) for arg in args]
-    return FunCall(fun=fun, args=[*make_node(args)])
+    return FunCall(fun=fun, args=[make_node(arg) for arg in args])
 
 
 # shift promotes its arguments to literals, therefore special
@@ -156,9 +156,7 @@ def make_node(o):
     if isinstance(o, common.Dimension):
         return AxisLiteral(value=o.value)
     if isinstance(o, tuple):
-        return tuple(make_node(arg) for arg in o)
-    if isinstance(o, list):
-        return list(make_node(arg) for arg in o)
+        return _f("make_tuple", *(make_node(arg) for arg in o))
     if o is None:
         return NoneLiteral()
     if hasattr(o, "fun"):
@@ -171,10 +169,7 @@ def trace_function_call(fun, *, args=None):
         args = (_s(param) for param in inspect.signature(fun).parameters.keys())
     body = fun(*list(args))
 
-    if isinstance(body, tuple):
-        return _f("make_tuple", *tuple(make_node(b) for b in body))
-    else:
-        return make_node(body) if body is not None else None
+    return make_node(body) if body is not None else None
 
 
 def lambdadef(fun):


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
index 3843b851..3f7a98f7 100644
--- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
+++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
@@ -84,15 +84,25 @@ def test_tuple_output(program_processor, stencil):
         assert np.allclose(inp2, out[1])
 
 
-def test_tuple_of_field_of_tuple_output(program_processor_no_gtfn_exec):
+@fundef
+def tuple_of_tuple_output1(inp1, inp2, inp3, inp4):
+    return (deref(inp1), deref(inp2)), (deref(inp3), deref(inp4))
+
+
+@fundef
+def tuple_of_tuple_output2(inp1, inp2, inp3, inp4):
+    return make_tuple(deref(inp1), deref(inp2)), make_tuple(deref(inp3), deref(inp4))
+
+
+@pytest.mark.parametrize(
+    "stencil",
+    [tuple_of_tuple_output1, tuple_of_tuple_output2],
+)
+def test_tuple_of_field_of_tuple_output(program_processor_no_gtfn_exec, stencil):
     program_processor, validate = program_processor_no_gtfn_exec
     if program_processor == run_dace_iterator:
         pytest.xfail("Not supported in DaCe backend: tuple returns")
 
-    @fundef
-    def stencil(inp1, inp2, inp3, inp4):
-        return make_tuple(deref(inp1), deref(inp2)), make_tuple(deref(inp3), deref(inp4))
-
     shape = [5, 7, 9]
     rng = np.random.default_rng()
     inp1 = gtx.np_as_located_field(IDim, JDim, KDim)(