From fa1c588f74e4706f39273aa453790d55d5d1df22 Mon Sep 17 00:00:00 2001 From: nfarabullini Date: Thu, 4 Jan 2024 10:07:53 +0100 Subject: [PATCH] reverted a couple of edits --- src/gt4py/next/embedded/nd_array_field.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index 8c48183024..622747b52b 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -161,13 +161,15 @@ def from_array( return cls(domain, array) def remap( - self: NdArrayField, connectivity: common.ConnectivityField + self: NdArrayField, connectivity: common.ConnectivityField | fbuiltins.FieldOffset ) -> NdArrayField: # For neighbor reductions, a FieldOffset is passed instead of an actual ConnectivityField if not common.is_connectivity_field(connectivity): + assert isinstance(connectivity, fbuiltins.FieldOffset) connectivity = connectivity.as_connectivity_field() assert common.is_connectivity_field(connectivity) + # Compute the new domain dim = connectivity.codomain dim_idx = self.domain.dim_index(dim)