[Update] Upgrade to the numpy 2x series

.github/workflows/fv3_translate_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   fv3_translate_tests:
-    uses: NOAA-GFDL/pyFV3/.github/workflows/translate.yaml@develop
+    uses: twicki/pyFV3/.github/workflows/translate.yaml@update/numpy_2x
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/pace_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   pace_main_tests:
-    uses: NOAA-GFDL/pace/.github/workflows/main_unit_tests.yaml@develop
+    uses: floriandeconinck/pace/.github/workflows/main_unit_tests.yaml@update/numpy_2x
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/shield_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   shield_translate_tests:
-    uses: NOAA-GFDL/pySHiELD/.github/workflows/translate.yaml@develop
+    uses: floriandeconinck/pySHiELD/.github/workflows/translate.yaml@update/numpy_2x
     with:
       component_trigger: true
       component_name: NDSL

ndsl/buffer.py

@@ -2,9 +2,10 @@
 
 import contextlib
 from collections.abc import Callable, Generator, Iterable
+from typing import Any
 
 import numpy as np
-from numpy.lib.index_tricks import IndexExpression
+import numpy.typing as npt
 
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.types import Allocator
@@ -16,7 +17,7 @@
 )
 
 
-BufferKey = tuple[Callable, Iterable[int], type]
+BufferKey = tuple[Callable, Iterable[int], npt.DTypeLike]
 BUFFER_CACHE: dict[BufferKey, list["Buffer"]] = {}
 
 
@@ -41,7 +42,7 @@ def __init__(self, key: BufferKey, array: np.ndarray):
 
     @classmethod
     def pop_from_cache(
-        cls, allocator: Allocator, shape: Iterable[int], dtype: type
+        cls, allocator: Allocator, shape: Iterable[int], dtype: npt.DTypeLike
     ) -> Buffer:
         """Retrieve or insert then retrieve of buffer from cache.
 
@@ -78,8 +79,8 @@ def finalize_memory_transfer(self) -> None:
     def assign_to(
         self,
         destination_array: np.ndarray,
-        buffer_slice: IndexExpression = np.index_exp[:],
-        buffer_reshape: IndexExpression = None,
+        buffer_slice: Any = np.index_exp[:],
+        buffer_reshape: Any | None = None,
     ) -> None:
         """Assign internal array to destination_array.
 
@@ -95,7 +96,7 @@ def assign_to(
             )
 
     def assign_from(
-        self, source_array: np.ndarray, buffer_slice: IndexExpression = np.index_exp[:]
+        self, source_array: np.ndarray, buffer_slice: Any = np.index_exp[:]
     ) -> None:
         """Assign source_array to internal array.
 
@@ -107,7 +108,7 @@ def assign_from(
 
 @contextlib.contextmanager
 def array_buffer(
-    allocator: Allocator, shape: Iterable[int], dtype: type
+    allocator: Allocator, shape: Iterable[int], dtype: npt.DTypeLike
 ) -> Generator[Buffer, Buffer, None]:
     """
     A context manager providing a contiguous array, which may be re-used between calls.
@@ -132,7 +133,7 @@ def send_buffer(
     allocator: Callable,
     array: np.ndarray,
     timer: Timer | None = None,
-) -> np.ndarray:
+) -> Generator[np.ndarray]:
     """A context manager ensuring that `array` is contiguous in a context where it is
     being sent as data, copying into a recycled buffer array if necessary.
 
@@ -166,7 +167,7 @@ def recv_buffer(
     allocator: Callable,
     array: np.ndarray,
     timer: Timer | None = None,
-) -> np.ndarray:
+) -> Generator[np.ndarray]:
     """A context manager ensuring that array is contiguous in a context where it is
     being used to receive data, using a recycled buffer array and then copying the
     result into array if necessary.

ndsl/comm/communicator.py

@@ -2,6 +2,7 @@
 
 import abc
 from collections.abc import Mapping, Sequence
+from types import ModuleType
 from typing import Any, Self, cast
 
 import numpy as np
@@ -16,7 +17,6 @@
 from ndsl.optional_imports import cupy
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.quantity import Quantity, QuantityHaloSpec, QuantityMetadata
-from ndsl.types import NumpyModule
 
 
 def to_numpy(array, dtype=None) -> np.ndarray:  # type: ignore[no-untyped-def]
@@ -83,7 +83,7 @@ def size(self) -> int:
         """Total number of ranks in this communicator"""
         return self.comm.Get_size()
 
-    def _maybe_force_cpu(self, module: NumpyModule) -> NumpyModule:
+    def _maybe_force_cpu(self, module: ModuleType) -> ModuleType:
         """
         Get a numpy-like module depending on configuration and
         Quantity original allocator.
@@ -223,7 +223,7 @@ def _get_gather_recv_quantity(
     ) -> Quantity:
         """Initialize a Quantity for use when receiving global data during gather"""
         recv_quantity = Quantity(
-            send_metadata.np.zeros(global_extent, dtype=send_metadata.dtype),  # type: ignore
+            send_metadata.np.zeros(global_extent, dtype=send_metadata.dtype),
             dims=send_metadata.dims,
             units=send_metadata.units,
             origin=tuple([0 for dim in send_metadata.dims]),
@@ -238,7 +238,7 @@ def _get_scatter_recv_quantity(
     ) -> Quantity:
         """Initialize a Quantity for use when receiving subtile data during scatter"""
         recv_quantity = Quantity(
-            send_metadata.np.zeros(shape, dtype=send_metadata.dtype),  # type: ignore
+            send_metadata.np.zeros(shape, dtype=send_metadata.dtype),
             dims=send_metadata.dims,
             units=send_metadata.units,
             backend=send_metadata.backend,
@@ -837,7 +837,7 @@ def _get_gather_recv_quantity(
         # needs to change the quantity dimensions since we add a "tile" dimension,
         # unlike for tile scatter/gather which retains the same dimensions
         recv_quantity = Quantity(
-            metadata.np.zeros(global_extent, dtype=metadata.dtype),  # type: ignore
+            metadata.np.zeros(global_extent, dtype=metadata.dtype),
             dims=(constants.TILE_DIM,) + metadata.dims,
             units=metadata.units,
             origin=(0,) + tuple([0 for dim in metadata.dims]),
@@ -859,7 +859,7 @@ def _get_scatter_recv_quantity(
         # needs to change the quantity dimensions since we remove a "tile" dimension,
         # unlike for tile scatter/gather which retains the same dimensions
         recv_quantity = Quantity(
-            metadata.np.zeros(shape, dtype=metadata.dtype),  # type: ignore
+            metadata.np.zeros(shape, dtype=metadata.dtype),
             dims=metadata.dims[1:],
             units=metadata.units,
             backend=metadata.backend,

ndsl/dsl/typing.py

@@ -1,6 +1,7 @@
 from typing import TypeAlias
 
 import numpy as np
+import numpy.typing as npt
 from gt4py.cartesian import gtscript
 
 from ndsl.dsl import NDSL_GLOBAL_PRECISION
@@ -110,7 +111,7 @@ def cast_to_index3d(val: tuple[int, ...]) -> Index3D:
     return val
 
 
-def is_float(dtype: type) -> bool:
+def is_float(dtype: npt.DTypeLike) -> bool:
     """Expected floating point type"""
     return dtype in [
         Float,

ndsl/halo/data_transformer.py

@@ -4,6 +4,8 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
+from types import ModuleType
+from typing import no_type_check
 from uuid import UUID, uuid1
 
 import numpy as np
@@ -22,7 +24,6 @@
 from ndsl.halo.rotate import rotate_scalar_data, rotate_vector_data
 from ndsl.optional_imports import cupy as cp
 from ndsl.quantity import Quantity, QuantityHaloSpec
-from ndsl.types import NumpyModule
 from ndsl.utils import device_synchronize
 
 
@@ -53,7 +54,11 @@ def _push_stream(stream: "cp.cuda.Stream") -> None:
 INDICES_CACHE: dict[str, "cp.ndarray"] = {}
 
 
-def _build_flatten_indices(  # type: ignore[no-untyped-def]
+# `array_value[...] = xxx` is failing mypy because of bad inference
+# of the type. We can't type ignore, because mypy also thinks that it
+# no needed (but if removed, it will fail...)
+@no_type_check
+def _build_flatten_indices(
     key,
     shape,
     slices: tuple[slice, ...],
@@ -186,7 +191,7 @@ class HaloDataTransformer(abc.ABC):
 
     def __init__(
         self,
-        np_module: NumpyModule,
+        np_module: ModuleType,
         exchange_descriptors_x: Sequence[HaloExchangeSpec],
         exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
     ) -> None:
@@ -237,7 +242,7 @@ def finalize(self) -> None:
 
     @staticmethod
     def get(
-        np_module: NumpyModule,
+        np_module: ModuleType,
         exchange_descriptors_x: Sequence[HaloExchangeSpec],
         exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
     ) -> HaloDataTransformer:
@@ -308,7 +313,7 @@ def _compile(self) -> None:
 
         # Compute required size
         buffer_size = 0
-        dtype = None
+        dtype = np.float32  # default that will be overriden or not used
         for edge_x in self._infos_x:
             buffer_size += edge_x.pack_buffer_size
             dtype = edge_x.specification.dtype
@@ -320,12 +325,12 @@ def _compile(self) -> None:
         self._pack_buffer = Buffer.pop_from_cache(
             self._np_module.zeros,
             (buffer_size,),
-            dtype,  # type: ignore[arg-type]
+            dtype,
         )
         self._unpack_buffer = Buffer.pop_from_cache(
             self._np_module.zeros,
             (buffer_size,),
-            dtype,  # type: ignore[arg-type]
+            dtype,
         )
 
     def ready(self) -> bool:
@@ -589,7 +594,7 @@ class _CuKernelArgs:
 
     def __init__(
         self,
-        np_module: NumpyModule,
+        np_module: ModuleType,
         exchange_descriptors_x: Sequence[HaloExchangeSpec],
         exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
     ) -> None:

ndsl/halo/updater.py

@@ -2,6 +2,7 @@
 
 from collections import defaultdict
 from collections.abc import Iterable, Mapping
+from types import ModuleType
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -14,7 +15,7 @@
 from ndsl.halo.rotate import rotate_scalar_data
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.quantity import Quantity, QuantityHaloSpec
-from ndsl.types import AsyncRequest, NumpyModule
+from ndsl.types import AsyncRequest
 from ndsl.utils import device_synchronize
 
 
@@ -95,7 +96,7 @@ def __del__(self) -> None:
     def from_scalar_specifications(
         cls,
         comm: Communicator,
-        numpy_like_module: NumpyModule,
+        numpy_like_module: ModuleType,
         specifications: Iterable[QuantityHaloSpec],
         boundaries: Iterable[Boundary],
         tag: int,
@@ -147,7 +148,7 @@ def from_scalar_specifications(
     def from_vector_specifications(
         cls,
         comm: Communicator,
-        numpy_like_module: NumpyModule,
+        numpy_like_module: ModuleType,
         specifications_x: Iterable[QuantityHaloSpec],
         specifications_y: Iterable[QuantityHaloSpec],
         boundaries: Iterable[Boundary],
@@ -475,7 +476,7 @@ def _Isend_vector_shared_boundary(
         ]
         return send_requests
 
-    def _maybe_force_cpu(self, module: NumpyModule) -> NumpyModule:
+    def _maybe_force_cpu(self, module: ModuleType) -> ModuleType:
         """
         Get a numpy-like module depending on configuration and
         Quantity original allocator.

ndsl/quantity/metadata.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import dataclasses
+from types import ModuleType
 from typing import Any
 
 import numpy as np
+import numpy.typing as npt
 
 from ndsl.config.backend import Backend
 from ndsl.optional_imports import cupy
-from ndsl.types import NumpyModule
 
 
 if cupy is None:
@@ -28,7 +29,7 @@ class QuantityMetadata:
     "Units of the quantity."
     data_type: type
     "ndarray-like type used to store the data."
-    dtype: type
+    dtype: npt.DTypeLike
     "dtype of the data in the ndarray-like object."
     backend: Backend
     "NDSL backend. Used for performance optimal data allocation."
@@ -39,7 +40,7 @@ def dim_lengths(self) -> dict[str, int]:
         return dict(zip(self.dims, self.extent))
 
     @property
-    def np(self) -> NumpyModule:
+    def np(self) -> ModuleType:
         """numpy-like module used to interact with the data."""
         if issubclass(self.data_type, cupy.ndarray):
             return cupy
@@ -72,5 +73,5 @@ class QuantityHaloSpec:
     origin: tuple[int, ...]
     extent: tuple[int, ...]
     dims: tuple[str, ...]
-    numpy_module: NumpyModule
+    numpy_module: ModuleType
     dtype: Any

ndsl/quantity/quantity.py

@@ -2,6 +2,7 @@
 
 import warnings
 from collections.abc import Iterable, Sequence
+from types import ModuleType
 from typing import Any, cast
 
 import dace
@@ -18,7 +19,6 @@
 from ndsl.optional_imports import cupy
 from ndsl.quantity.bounds import BoundedArrayView
 from ndsl.quantity.metadata import QuantityHaloSpec, QuantityMetadata
-from ndsl.types import NumpyModule
 
 
 if cupy is None:
@@ -326,7 +326,7 @@ def data_as_xarray(self) -> xr.DataArray:
         return xr.DataArray(data, dims=self.dims, attrs=self.attrs)
 
     @property
-    def np(self) -> NumpyModule:
+    def np(self) -> ModuleType:
         return self.metadata.np
 
     @property
@@ -408,7 +408,7 @@ def transpose(
         target_dims = _collapse_dims(target_dims, self.dims)
         transpose_order = [self.dims.index(dim) for dim in target_dims]
         transposed = Quantity(
-            self.np.transpose(self.data, transpose_order),  # type: ignore[attr-defined]
+            self.np.transpose(self.data, transpose_order),
             dims=_transpose_sequence(self.dims, transpose_order),
             units=self.units,
             origin=_transpose_sequence(self.origin, transpose_order),