From 41a9f74be752cf254312e68dfaf1045090c6f763 Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 8 Nov 2019 18:42:13 -0500 Subject: [PATCH 1/3] Allow CuPy arrays through `ensure_ndarray` --- numcodecs/compat.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/numcodecs/compat.py b/numcodecs/compat.py index 52827ca6..4809e5df 100644 --- a/numcodecs/compat.py +++ b/numcodecs/compat.py @@ -6,6 +6,7 @@ import array +import cupy as cp import numpy as np @@ -48,8 +49,8 @@ def ensure_ndarray(buf): """ - if isinstance(buf, np.ndarray): - # already a numpy array + if isinstance(buf, (np.ndarray, cp.ndarray)): + # already an array arr = buf elif isinstance(buf, array.array) and buf.typecode in 'cu': From 12337d4defc8e395af4cd9243c8268ffd7eafcaf Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 8 Nov 2019 18:47:14 -0500 Subject: [PATCH 2/3] Add a workaround for CuPy arrays in `ensure_bytes` As CuPy arrays (currently) do not have a `tobytes` method, workaround this in `ensure_bytes` by coercing CuPy arrays back to NumPy arrays before calling `tobytes`. This way it can just use NumPy's `tobytes` method. Should be unneeded in a future version of CuPy. --- numcodecs/compat.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/numcodecs/compat.py b/numcodecs/compat.py index 4809e5df..10b034f8 100644 --- a/numcodecs/compat.py +++ b/numcodecs/compat.py @@ -152,6 +152,12 @@ def ensure_bytes(buf): if arr.dtype == object: raise TypeError('object arrays are not supported') + # Force CuPy arrays to NumPy arrays + # because they don't have a `tobytes` method (yet) + # xref: https://github.com/cupy/cupy/pull/2617 + if isinstance(arr, cp.ndarray): + arr = arr.get() + # create bytes buf = arr.tobytes(order='A') From d86fe9116ffd4ff7f90d353af3a9b2c34303230b Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 8 Nov 2019 18:56:17 -0500 Subject: [PATCH 3/3] Force CuPy arrays to NumPy arrays to decode text As the decoding algorithm requires something that supports the buffer protocol and assumes it is working on host memory, this won't work on a CuPy array. So convert the CuPy array to a NumPy array first so that it can proceed through the usual decoding path. --- numcodecs/compat.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/numcodecs/compat.py b/numcodecs/compat.py index 10b034f8..a239e838 100644 --- a/numcodecs/compat.py +++ b/numcodecs/compat.py @@ -167,6 +167,12 @@ def ensure_bytes(buf): def ensure_text(s, encoding='utf-8'): if not isinstance(s, text_type): s = ensure_contiguous_ndarray(s) + + # Force CuPy arrays to NumPy arrays + # as they support the buffer protocol + if isinstance(s, cp.ndarray): + s = s.get() + s = codecs.decode(s, encoding) return s