Basic setitem and getitem in linked array cupy

xsuite · Mar 30, 2022 · 680c4ad · 680c4ad
1 parent ffbb46a
commit 680c4ad
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py
@@ -53,7 +53,17 @@ def _as_cupy(self):
                                    strides=self.strides)
 
         def _basic_setitem(self, indx, val):
-            cupy.ndarray.__setitem__(self._as_cupy, indx, val)
+            if hasattr(val, '_as_cupy'):
+                val = val._as_cupy()
+            if hasattr(indx, '_as_cupy'):
+                indx = indx._as_cupy()
+            cupy.ndarray.__setitem__(self._as_cupy(), indx, val)
+
+        def __getitem__(self, indx):
+            if hasattr(indx, '_as_cupy'):
+                indx = indx._as_cupy()
+            return cupy.ndarray.__getitem__(self._as_cupy(), indx)
+
 
 #        # The following methods are generated by this code
 #        binops = [