[prototype] Approach 3: PartialTypeSpec lazy generic types (vs Stage-0)

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

@@ -831,7 +831,7 @@ def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> foast.Call:
             print(f"Warning: return type of '{func_name}' might be inconsistent (not implemented).")
 
         # deduce return type
-        return_type: Optional[ts.FieldType | ts.ScalarType | ts.TypeVarType] = None
+        return_type: Optional[ts.FieldType | ts.ScalarType] = None
         if (
             func_name
             in fbuiltins.UNARY_MATH_NUMBER_BUILTIN_NAMES + fbuiltins.UNARY_MATH_FP_BUILTIN_NAMES

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -181,16 +181,12 @@ def _transform_by_pattern(
                 lambda x: next(uids["__tmp"]),
                 result_collection_constructor=as_tuple,
             )(tmp_expr.type)
-            # the lowered IR is concrete, so `extract_dtype` never yields a `TypeVarType` here
             tmp_dtypes: (
                 ts.ScalarType | ts.ListType | tuple[ts.ScalarType | ts.ListType | tuple, ...]
-            ) = cast(
-                "ts.ScalarType | ts.ListType | tuple[ts.ScalarType | ts.ListType | tuple, ...]",
-                type_info.tree_map_type(
-                    type_info.extract_dtype,
-                    result_collection_constructor=as_tuple,
-                )(tmp_expr.type),
-            )
+            ) = type_info.tree_map_type(
+                type_info.extract_dtype,
+                result_collection_constructor=as_tuple,
+            )(tmp_expr.type)
 
             tmp_domains: SymbolicDomain | tuple[SymbolicDomain | tuple, ...] = tmp_expr.annex.domain
 

src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py

@@ -104,9 +104,7 @@ class MemletExpr:
 
     @property
     def gt_dtype(self) -> ts.ScalarType | ts.ListType:
-        dtype = self.gt_field.dtype
-        assert isinstance(dtype, (ts.ScalarType, ts.ListType))
-        return dtype
+        return self.gt_field.dtype
 
     def __post_init__(self) -> None:
         if isinstance(self.gt_dtype, ts.ListType):

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py

@@ -896,7 +896,6 @@ def _add_storage(
                 dc_dtype = gtx_dace_args.as_dace_type(gt_type.dtype)
                 all_dims = gt_type.dims
             else:  # for 'ts.ListType' use 'offset_type' as local dimension
-                assert isinstance(gt_type.dtype, ts.ListType)
                 assert gt_type.dtype.offset_type is not None
                 assert gt_type.dtype.offset_type.kind == gtx_common.DimensionKind.LOCAL
                 assert isinstance(gt_type.dtype.element_type, ts.ScalarType)

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_concat_where.py

@@ -143,7 +143,6 @@ def testee(interior: cases.IJKField, boundary: cases.IJField) -> cases.IJKField:
     if isinstance(output_type.dtype, ts.ScalarType):
         all_dims = gtx_common.order_dimensions(output_type.dims)
     else:
-        assert isinstance(output_type.dtype, ts.ListType)
         assert output_type.dtype.offset_type
         all_dims = gtx_common.order_dimensions([*output_type.dims, output_type.dtype.offset_type])
 

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py

@@ -622,9 +622,7 @@ def _handle_dataflow_result_of_nested_sdfg(
         None,
         dace.Memlet.from_array(outer_dataname, outer_desc),
     )
-    output_dtype = inner_data.gt_type.dtype
-    assert isinstance(output_dtype, (ts.ScalarType, ts.ListType))
-    output_expr = gtir_dataflow.ValueExpr(outer_node, output_dtype)
+    output_expr = gtir_dataflow.ValueExpr(outer_node, inner_data.gt_type.dtype)
     return gtir_dataflow.DataflowOutputEdge(outer_ctx.state, output_expr)
 
 

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_types.py

@@ -86,11 +86,9 @@ def get_local_view(
                 (dim, dace.symbolic.SymExpr(0) if self.origin is None else self.origin[i])
                 for i, dim in enumerate(self.gt_type.dims)
             ]
-            dtype = self.gt_type.dtype
-            # `FieldType.dtype` is widened to include `TypeVarType`, but generic operators are
-            # monomorphized before lowering, so only concrete dtypes reach here.
-            assert isinstance(dtype, (ts.ScalarType, ts.ListType))
-            return gtir_dataflow.IteratorExpr(self.dc_node, dtype, field_origin, it_indices)
+            return gtir_dataflow.IteratorExpr(
+                self.dc_node, self.gt_type.dtype, field_origin, it_indices
+            )
 
         raise NotImplementedError(f"Node type {type(self.gt_type)} not supported.")
 

src/gt4py/next/type_system/_approach3_deduction_sketch.py

@@ -0,0 +1,102 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""PROTOTYPE-ONLY: the FOAST type-deduction impact of Approach 3 (`PartialTypeSpec`).
+
+This is *not* wired into the toolchain. It is a faithful, compilable sketch of how the three
+representative deduction operations in `ffront/foast_passes/type_deduction.py` would have to grow
+a `PartialTypeSpec` branch, so the frontend invasiveness can be judged from real code instead of
+prose.
+
+The contrast with Stage-0 (TypeVar-inside-FieldType): there, a generic field IS a `FieldType`
+(with `dtype=TypeVarType`), so `case ts.FieldType(dims, dtype)` and `ts.FieldType(...)`
+construction "just work" and the *predicates* (`promote`, `extract_dtype`, ...) were taught about
+`TypeVarType` once, centrally. With `PartialTypeSpec` a generic field is a different object, so
+every site that *matches* or *constructs* a `FieldType` in a path a generic type can reach needs a
+parallel branch.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from gt4py.next.type_system import partial_type_info as pti, type_specifications as ts
+
+
+# --- representative site 1: `with_altered_scalar_kind` (ffront/foast_passes/type_deduction.py) ---
+# Stage-0 today (works unchanged on a generic field, because dtype just happens to be a TypeVarType):
+#
+#     if isinstance(type_spec, ts.FieldType):
+#         return ts.FieldType(dims=type_spec.dims,
+#                             dtype=with_altered_scalar_kind(type_spec.dtype, new_scalar_kind))
+#
+# Approach 3: the predicate flips the dtype to bool, which is concrete, so the result is a concrete
+# field even though the input was generic. But the *input* is now a PartialTypeSpec, so we must
+# match it explicitly and reach into `fields`:
+def with_altered_scalar_kind(
+    type_spec: ts.TypeSpec, new_scalar_kind: ts.ScalarKind
+) -> ts.ScalarType | ts.FieldType:
+    if isinstance(type_spec, ts.PartialTypeSpec) and type_spec.target is ts.FieldType:
+        fields = dict(type_spec.fields)
+        # a generic dtype becomes a concrete `bool` scalar here -> a concrete FieldType
+        return ts.FieldType(dims=fields["dims"], dtype=ts.ScalarType(kind=new_scalar_kind))
+    if isinstance(type_spec, ts.FieldType):
+        return ts.FieldType(dims=type_spec.dims, dtype=ts.ScalarType(kind=new_scalar_kind))
+    if isinstance(type_spec, ts.ScalarType):
+        return ts.ScalarType(kind=new_scalar_kind, shape=type_spec.shape)
+    raise ValueError(f"Expected field or scalar type, got '{type_spec}'.")
+
+
+# --- representative site 2: binary-op promotion (`_deduce_binop_type` -> type_info.promote) ---
+# Stage-0 today: `type_info.promote(left.type, right.type)` — promote() was taught that
+# `promote(T, T) = T` and that mixing T with a concrete dtype is the D3 strict-no-promotion error.
+#
+# Approach 3: `type_info.promote` is concrete-only again. Generic promotion is a *separate* code
+# path keyed on "either operand is a PartialTypeSpec". The D3 rule lives here now, not in promote.
+def promote_for_binop(left: ts.TypeSpec, right: ts.TypeSpec) -> ts.TypeSpec:
+    l_gen, r_gen = isinstance(left, ts.PartialTypeSpec), isinstance(right, ts.PartialTypeSpec)
+    if l_gen or r_gen:
+        if not (l_gen and r_gen):
+            # D3: a generic dtype combined with a concrete one is a (decoration-time) error.
+            raise ValueError(
+                f"Could not promote '{left}' and '{right}': a generic dtype can only be"
+                " combined with the same type variable."
+            )
+        l_tvars, r_tvars = pti.type_vars(left), pti.type_vars(right)
+        if set(l_tvars) != set(r_tvars):
+            raise ValueError(f"Could not promote distinct type variables '{left}' and '{right}'.")
+        # `promote(T, T) = T`: dims still promote like concrete fields; here we just return one.
+        # (A real impl would promote the `dims` field and keep the shared TypeVar dtype.)
+        return left
+    # both concrete -> delegate to the untouched concrete promote
+    from gt4py.next.type_system import type_info
+
+    return type_info.promote(left, right)  # type: ignore[arg-type]  # concrete branch
+
+
+# --- representative site 3: subscript / dimension drop (`visit_Subscript`) ---
+# Stage-0 today:
+#     case ts.FieldType(dims=dims, dtype=dtype):
+#         new_type = ts.FieldType(dims=[d for d in dims if d != idx.dim], dtype=dtype)
+#
+# Approach 3: add a parallel `PartialTypeSpec` arm that rebuilds the partial with the same dtype
+# (still a TypeVar) but reduced dims:
+def subscript_drop_dim(value_type: ts.TypeSpec, dropped: object) -> Optional[ts.TypeSpec]:
+    match value_type:
+        case ts.PartialTypeSpec(target=target, fields=fields) if target is ts.FieldType:
+            f = dict(fields)
+            # `PartialField = object`, so the field value is untyped and must be re-narrowed.
+            dims = f["dims"]
+            assert isinstance(dims, list)
+            return ts.PartialTypeSpec(
+                target=ts.FieldType,
+                fields=(("dims", [d for d in dims if d != dropped]), ("dtype", f["dtype"])),
+            )
+        case ts.FieldType(dims=dims, dtype=dtype):
+            return ts.FieldType(dims=[d for d in dims if d != dropped], dtype=dtype)
+    return None

src/gt4py/next/type_system/partial_type_info.py

@@ -0,0 +1,156 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Frontend-only utilities for `ts.PartialTypeSpec` (Approach 3 prototype).
+
+Genericity is expressed by wrapping a target `TypeSpec` class together with its constructor
+fields, where some fields are `TypeVarType` placeholders. None of this is visible downstream of
+specialization, so the IR/backend never has to narrow `FieldType.dtype`.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from gt4py.eve import extended_typing as xtyping
+from gt4py.next.type_system import type_specifications as ts
+
+
+def _walk_typevars(value: Any) -> xtyping.Iterator[ts.TypeVarType]:
+    """Yield every `TypeVarType` reachable inside a (possibly partial) field value."""
+    if isinstance(value, ts.TypeVarType):
+        yield value
+    elif isinstance(value, ts.PartialTypeSpec):
+        for _, v in value.fields:
+            yield from _walk_typevars(v)
+    elif isinstance(value, (list, tuple)):
+        for v in value:
+            yield from _walk_typevars(v)
+    elif isinstance(value, ts.TupleType):
+        for v in value.types:
+            yield from _walk_typevars(v)
+    elif isinstance(value, ts.FieldType):
+        yield from _walk_typevars(value.dtype)
+    elif isinstance(value, ts.ListType):
+        yield from _walk_typevars(value.element_type)
+
+
+def is_generic(type_: Any) -> bool:
+    """A `PartialTypeSpec` (or a structure containing one / a bare `TypeVarType`) is generic."""
+    return next(_walk_typevars(type_), None) is not None
+
+
+def type_vars(type_: Any) -> dict[str, ts.TypeVarType]:
+    """Collect the type variables occurring in ``type_``, keyed by name."""
+    return {var.name: var for var in _walk_typevars(type_)}
+
+
+def _substitute_value(value: Any, binding: xtyping.Mapping[str, ts.ScalarType]) -> Any:
+    """Replace bound type variables inside a single field value, recursing into nested partials."""
+    if isinstance(value, ts.TypeVarType):
+        return binding.get(value.name, value)
+    if isinstance(value, ts.PartialTypeSpec):
+        return substitute(value, binding)
+    if isinstance(value, list):
+        return [_substitute_value(v, binding) for v in value]
+    if isinstance(value, tuple):
+        return tuple(_substitute_value(v, binding) for v in value)
+    return value
+
+
+def substitute(
+    partial: ts.PartialTypeSpec, binding: xtyping.Mapping[str, ts.ScalarType]
+) -> ts.PartialTypeSpec:
+    """Return a `PartialTypeSpec` with the bound type variables substituted (still possibly partial)."""
+    return ts.PartialTypeSpec(
+        target=partial.target,
+        fields=tuple((name, _substitute_value(value, binding)) for name, value in partial.fields),
+    )
+
+
+def _materialize(value: Any) -> Any:
+    """Turn any nested `PartialTypeSpec` into its concrete `target`, recursively.
+
+    Precondition: ``value`` contains no unbound type variables.
+    """
+    if isinstance(value, ts.PartialTypeSpec):
+        return value.target(**{name: _materialize(v) for name, v in value.fields})
+    if isinstance(value, list):
+        return [_materialize(v) for v in value]
+    if isinstance(value, tuple):
+        return tuple(_materialize(v) for v in value)
+    return value
+
+
+def specialize(partial: ts.PartialTypeSpec, binding: xtyping.Mapping[str, ts.ScalarType]) -> ts.TypeSpec:
+    """Substitute all type variables and instantiate the concrete ``target`` `TypeSpec`.
+
+    Raises if a type variable remains unbound after substitution (the partial is still generic).
+    Nested `PartialTypeSpec`s are materialized too, so the result is fully concrete.
+    """
+    substituted = substitute(partial, binding)
+    if is_generic(substituted):
+        unbound = ", ".join(sorted(type_vars(substituted)))
+        raise ValueError(f"Cannot specialize '{partial}': unbound type variable(s) {unbound}.")
+    result = _materialize(substituted)
+    assert isinstance(result, ts.TypeSpec)
+    return result
+
+
+def _bind_var(var: ts.TypeVarType, dtype: ts.TypeSpec) -> dict[str, ts.ScalarType]:
+    if not isinstance(dtype, ts.ScalarType):
+        return {}
+    if dtype not in var.constraints:
+        raise ValueError(f"'{dtype}' does not satisfy the constraints of type variable '{var}'.")
+    return {var.name: dtype}
+
+
+def _merge(parts: xtyping.Iterable[dict[str, ts.ScalarType]]) -> dict[str, ts.ScalarType]:
+    binding: dict[str, ts.ScalarType] = {}
+    for part in parts:
+        for name, dtype in part.items():
+            if (prev := binding.get(name)) is not None and prev != dtype:
+                raise ValueError(
+                    f"Type variable '{name}' is bound inconsistently: '{prev}' and '{dtype}'."
+                )
+            binding[name] = dtype
+    return binding
+
+
+def _bind_value(param_value: Any, arg: ts.TypeSpec) -> dict[str, ts.ScalarType]:
+    """Bind type variables in one partial field value against a concrete argument type."""
+    if isinstance(param_value, ts.TypeVarType):
+        return _bind_var(param_value, arg)
+    if isinstance(param_value, ts.PartialTypeSpec):
+        return bind(param_value, arg)
+    return {}
+
+
+def bind(partial: ts.PartialTypeSpec, arg: ts.TypeSpec) -> dict[str, ts.ScalarType]:
+    """Bind the type variables in ``partial`` by structurally matching the concrete ``arg``.
+
+    Only the dtype field participates today (dtype-scoped, per ADR 0024); generalizing to other
+    fields (e.g. dims) is a matter of widening the structural match below.
+    """
+    if partial.target is ts.FieldType:
+        dtype_value = dict(partial.fields).get("dtype")
+        # scalar arguments are promoted to zero-dimensional fields
+        arg_dtype = arg.dtype if isinstance(arg, ts.FieldType) else arg
+        return _bind_value(dtype_value, arg_dtype)
+    return {}
+
+
+def bind_many(
+    params: Sequence[ts.TypeSpec], args: Sequence[ts.TypeSpec]
+) -> dict[str, ts.ScalarType]:
+    """Compute the binding for a full parameter/argument list (mixed partial and concrete)."""
+    return _merge(
+        bind(param, arg)
+        for param, arg in zip(params, args)
+        if isinstance(param, ts.PartialTypeSpec)
+    )