Add top-level astype function for Dask

sgkit-dev · Sep 19, 2024 · d522c0f · d522c0f
1 parent c359718
commit d522c0f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/sgkit/distarray.py b/sgkit/distarray.py
@@ -8,3 +8,9 @@
 else:
     # default to dask
     from dask.array import *  # noqa: F401, F403
+
+    # dask doesn't have a top-level astype required by the array API
+    def astype(x, dtype, /, *, copy=True):  # pragma: no cover
+        if not copy and dtype == x.dtype:
+            return x
+        return x.astype(dtype=dtype, copy=copy)