From 28edadb2b62306cce2ba862eebbb420d1c357ee5 Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Mon, 16 Sep 2024 13:04:42 +0100
Subject: [PATCH] Add top-level astype function for Dask

---
 sgkit/distarray.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/sgkit/distarray.py b/sgkit/distarray.py
index fe3bda79..5900334e 100644
--- 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):
+        if not copy and dtype == x.dtype:
+            return x
+        return x.astype(dtype=dtype, copy=copy)