modified assertion

GridTools · Jan 31, 2024 · 7d4e8c3 · 7d4e8c3
1 parent 7fb99e2
commit 7d4e8c3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py
@@ -207,8 +207,8 @@ def restrict(self, index: common.AnyIndexSpec) -> common.Field:
         new_domain, buffer_slice = self._slice(index)
         new_buffer = self.ndarray[buffer_slice]
         if new_domain.ndim == 0:
-            assert core_defs.is_scalar_type(new_buffer)
-            return self._scalar_to_field(new_buffer)  # type: ignore[return-value]
+            assert issubclass(self.array_ns.dtype(type(new_buffer)), core_defs.SCALAR_TYPES)
+            return self._scalar_to_field(new_buffer)  # type: ignore[return-value, arg-type]
         return self.__class__.from_array(new_buffer, domain=new_domain)
 
     __getitem__ = restrict