Add top-level astype function for Dask

sgkit-dev · Sep 16, 2024 · 4588f19 · 4588f19
1 parent e8a7890
commit 4588f19
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)