linked array copy return nparray on cpu

xobjects/_version.py

@@ -1 +1 @@
-__version__ = "0.4.1"
+__version__ = "0.4.3"

xobjects/context.py

@@ -210,19 +210,30 @@ class XContext(ABC):
 
     def __init__(self):
         self._kernels = KernelDict()
-        self._buffers = []
+        self._buffers = weakref.WeakSet()
+        self._allocations = 0
 
     def __str__(self):
         return type(self).__name__
 
     def new_buffer(self, capacity=1048576):
         buf = self._make_buffer(capacity=capacity)
-        self.buffers.append(weakref.finalize(buf, log.debug, "free buf"))
+        self._buffers.add(buf)
+        self._allocations += 1
         return buf
 
+    def __getstate__(self):
+        state = self.__dict__
+        del state["_buffers"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
+
     @property
     def buffers(self):
-        return self._buffers
+        return list(self._buffers)
 
     @property
     def kernels(self):

xobjects/context_cpu.py

@@ -10,6 +10,7 @@
 import uuid
 from pathlib import Path
 from typing import Callable, Dict, List, Sequence, Tuple
+import weakref
 
 from .general import _print
 
@@ -106,6 +107,9 @@ def _build_view(cls, a):
             order="C",
         )
 
+    def copy(self):
+        return np.array(self)
+
 
 def _so_for_module_name(name, containing_dir=".") -> Path:
     # The so file name is something like:
@@ -642,10 +646,12 @@ def openmp_enabled(self):
     def __getstate__(self):
         state = self.__dict__.copy()
         state["_kernels"] = {}
+        del state["_buffers"]
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
 
 
 class BufferByteArray(XBuffer):

xobjects/context_cupy.py

@@ -473,6 +473,9 @@ def build_kernels(
 
         return out_kernels
 
+    def __str__(self):
+        return f"{type(self).__name__}:{cupy.cuda.get_device_id()}"
+
     def nparray_to_context_array(self, arr):
         """
         Copies a numpy array to the device memory.
@@ -631,7 +634,6 @@ class KernelCupy(object):
     def __init__(
         self, function, description, block_size, context, shared_mem_size_bytes
     ):
-
         self.function = function
         self.description = description
         self.block_size = block_size

xobjects/context_pyopencl.py

@@ -218,7 +218,9 @@ def build_kernels(
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        prg = cl.Program(self.context, specialized_source).build()
+        prg = cl.Program(self.context, specialized_source).build(
+            options="-cl-std=CL2.0",
+        )
 
         out_kernels = {}
         for pyname, kernel in kernel_descriptions.items():
@@ -236,6 +238,11 @@ def build_kernels(
 
         return out_kernels
 
+    def __str__(self):
+        platform_id = cl.get_platforms().index(self.platform)
+        device_id = self.platform.get_devices().index(self.device)
+        return f"{type(self).__name__}:{platform_id}.{device_id}"
+
     def nparray_to_context_array(self, arr):
         """
         Copies a numpy array to the device memory.

xobjects/general.py

@@ -3,6 +3,7 @@
 # Copyright (c) CERN, 2024.                   #
 # ########################################### #
 from numpy.testing import assert_allclose as np_assert_allclose
+import numpy as np
 
 
 class Print:
@@ -21,4 +22,12 @@ def assert_allclose(a, b, rtol=1e-7, atol=1e-7):
         a = a.get()
     if hasattr(b, "get"):
         b = b.get()
+    try:
+        a = np.squeeze(a)
+    except:
+        pass
+    try:
+        b = np.squeeze(b)
+    except:
+        pass
     np_assert_allclose(a, b, rtol=rtol, atol=atol)