feat[next]: Multiple output domains

src/gt4py/next/embedded/operators.py

@@ -11,6 +11,7 @@
 
 from gt4py import eve
 from gt4py._core import definitions as core_defs
+from gt4py.eve import extended_typing as xtyping
 from gt4py.next import common, errors, field_utils, utils
 from gt4py.next.embedded import common as embedded_common, context as embedded_context
 from gt4py.next.field_utils import get_array_ns
@@ -108,7 +109,9 @@ def field_operator_call(op: EmbeddedOperator[_R, _P], args: Any, kwargs: Any) ->
 
         domain = kwargs.pop("domain", None)
 
-        out_domain = common.domain(domain) if domain is not None else _get_out_domain(out)
+        out_domain = (
+            utils.tree_map(common.domain)(domain) if domain is not None else _get_out_domain(out)
+        )
 
         new_context_kwargs["closure_column_range"] = _get_vertical_range(out_domain)
 
@@ -128,6 +131,7 @@ def field_operator_call(op: EmbeddedOperator[_R, _P], args: Any, kwargs: Any) ->
         return op(*args, **kwargs)
 
 
+@utils.tree_map
 def _get_vertical_range(domain: common.Domain) -> common.NamedRange | eve.NothingType:
     vertical_dim_filtered = [nr for nr in domain if nr.dim.kind == common.DimensionKind.VERTICAL]
     assert len(vertical_dim_filtered) <= 1
@@ -137,17 +141,19 @@ def _get_vertical_range(domain: common.Domain) -> common.NamedRange | eve.Nothin
 def _tuple_assign_field(
     target: tuple[common.MutableField | tuple, ...] | common.MutableField,
     source: tuple[common.Field | tuple, ...] | common.Field,
-    domain: common.Domain,
+    domain: xtyping.MaybeNestedInTuple[common.Domain],
 ) -> None:
     @utils.tree_map
-    def impl(target: common.MutableField, source: common.Field) -> None:
+    def impl(target: common.MutableField, source: common.Field, domain: common.Domain) -> None:
         if isinstance(source, common.Field):
             target[domain] = source[domain]
         else:
             assert core_defs.is_scalar_type(source)
             target[domain] = source
 
-    impl(target, source)
+    if not isinstance(domain, tuple):
+        domain = utils.tree_map(lambda _: domain)(target)
+    impl(target, source, domain)
 
 
 def _intersect_scan_args(

src/gt4py/next/ffront/decorator.py

@@ -714,8 +714,10 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
                 raise errors.MissingArgumentError(None, "out", True)
             out = kwargs.pop("out")
             if "domain" in kwargs:
-                domain = common.domain(kwargs.pop("domain"))
-                out = utils.tree_map(lambda f: f[domain])(out)
+                domain = utils.tree_map(common.domain)(kwargs.pop("domain"))
+                if not isinstance(domain, tuple):
+                    domain = utils.tree_map(lambda _: domain)(out)
+                out = utils.tree_map(lambda f, dom: f[dom])(out, domain)
 
             args, kwargs = type_info.canonicalize_arguments(
                 self.foast_stage.foast_node.type, args, kwargs

src/gt4py/next/ffront/past_passes/type_deduction.py

@@ -37,6 +37,53 @@ def _is_integral_scalar(expr: past.Expr) -> bool:
     return isinstance(expr.type, ts.ScalarType) and type_info.is_integral(expr.type)
 
 
+def _validate_domain_out(
+    dom: past.Dict | past.TupleExpr,
+    out: ts.TypeSpec,
+    is_nested: bool = False,
+) -> None:
+    if isinstance(dom, past.Dict):
+        # Only reject tuple outputs if nested
+        if is_nested and isinstance(out, ts.TupleType):
+            raise ValueError("Domain dict cannot map to tuple outputs.")
+        assert not (is_nested and isinstance(out, past.TupleExpr))
+
+        if len(dom.keys_) == 0:
+            raise ValueError("Empty domain not allowed.")
+
+        for dim in dom.keys_:
+            if not isinstance(dim.type, ts.DimensionType):
+                raise ValueError(
+                    f"Only 'Dimension' allowed in domain dictionary keys, got '{dim}' which is of type '{dim.type}'."
+                )
+
+        for domain_values in dom.values_:
+            if len(domain_values.elts) != 2:
+                raise ValueError(
+                    f"Only 2 values allowed in domain range, got {len(domain_values.elts)}."
+                )
+            if any(not _is_integral_scalar(el) for el in domain_values.elts):
+                raise ValueError(
+                    f"Only integer values allowed in domain range, got '{domain_values.elts[0].type}' and '{domain_values.elts[1].type}'."
+                )
+
+    elif isinstance(dom, past.TupleExpr):
+        if isinstance(out, ts.TupleType):
+            out_elts = out.types
+        else:
+            raise ValueError(f"Tuple domain requires tuple output, got {type(out)}.")
+
+        if len(dom.elts) != len(out_elts):
+            raise ValueError("Mismatched tuple lengths between domain and output.")
+
+        for d, o in zip(dom.elts, out_elts, strict=True):
+            assert isinstance(d, (past.Dict, past.TupleExpr))
+            _validate_domain_out(d, o, is_nested=True)
+
+    else:
+        raise ValueError(f"'domain' must be Dict or TupleExpr, got {type(dom)}.")
+
+
 def _validate_operator_call(new_func: past.Name, new_kwargs: dict) -> None:
     """
     Perform checks for domain and output field types.
@@ -53,32 +100,11 @@ def _validate_operator_call(new_func: past.Name, new_kwargs: dict) -> None:
 
     if "out" not in new_kwargs:
         raise ValueError("Missing required keyword argument 'out'.")
-    if "domain" in new_kwargs:
+    if (domain := new_kwargs.get("domain")) is not None:
         _ensure_no_sliced_field(new_kwargs["out"])
-
-        domain_kwarg = new_kwargs["domain"]
-        if not isinstance(domain_kwarg, past.Dict):
-            raise ValueError(f"Only Dictionaries allowed in 'domain', got '{type(domain_kwarg)}'.")
-
-        if len(domain_kwarg.values_) == 0 and len(domain_kwarg.keys_) == 0:
-            raise ValueError("Empty domain not allowed.")
-
-        for dim in domain_kwarg.keys_:
-            if not isinstance(dim.type, ts.DimensionType):
-                raise ValueError(
-                    f"Only 'Dimension' allowed in domain dictionary keys, got '{dim}' which is of type '{dim.type}'."
-                )
-        for domain_values in domain_kwarg.values_:
-            if len(domain_values.elts) != 2:
-                raise ValueError(
-                    f"Only 2 values allowed in domain range, got {len(domain_values.elts)}."
-                )
-            if not _is_integral_scalar(domain_values.elts[0]) or not _is_integral_scalar(
-                domain_values.elts[1]
-            ):
-                raise ValueError(
-                    f"Only integer values allowed in domain range, got '{domain_values.elts[0].type}' and '{domain_values.elts[1].type}'."
-                )
+        out = new_kwargs["out"]
+        assert isinstance(out, past.Expr) and out.type is not None
+        _validate_domain_out(domain, out.type)
 
 
 class ProgramTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTranslator):
@@ -131,11 +157,22 @@ def visit_Attribute(self, node: past.Attribute, **kwargs: Any) -> past.Attribute
             type=getattr(new_value.type, node.attr),
         )
 
+    def visit_Dict(self, node: past.Dict, **kwargs: Any) -> past.Dict:
+        # the only supported dict for now is in domain specification
+        keys = self.visit(node.keys_, **kwargs)
+        assert all(isinstance(key.type, ts.DimensionType) for key in keys)
+        return past.Dict(
+            keys_=keys,
+            values_=self.visit(node.values_, **kwargs),
+            location=node.location,
+            type=ts.DomainType(dims=[key.type.dim for key in keys]),
+        )
+
     def visit_TupleExpr(self, node: past.TupleExpr, **kwargs: Any) -> past.TupleExpr:
         elts = self.visit(node.elts, **kwargs)
-        return past.TupleExpr(
-            elts=elts, type=ts.TupleType(types=[el.type for el in elts]), location=node.location
-        )
+        ttype = ts.TupleType(types=[elt.type for elt in elts])
+
+        return past.TupleExpr(elts=elts, type=ttype, location=node.location)
 
     def _deduce_binop_type(
         self, node: past.BinOp, *, left: past.Expr, right: past.Expr, **kwargs: Any