Convert array interface to numpy or cupy array
(interface: ArrayInf, zero_copy: bool = False)
| 121 | |
| 122 | |
| 123 | def from_array_interface(interface: ArrayInf, zero_copy: bool = False) -> NumpyOrCupy: |
| 124 | """Convert array interface to numpy or cupy array""" |
| 125 | |
| 126 | class Array: |
| 127 | """Wrapper type for communicating with numpy and cupy.""" |
| 128 | |
| 129 | _interface: Optional[ArrayInf] = None |
| 130 | |
| 131 | @property |
| 132 | def __array_interface__(self) -> Optional[ArrayInf]: |
| 133 | return self._interface |
| 134 | |
| 135 | @__array_interface__.setter |
| 136 | def __array_interface__(self, interface: ArrayInf) -> None: |
| 137 | self._interface = copy.copy(interface) |
| 138 | # Convert some fields to tuple as required by numpy |
| 139 | self._interface["shape"] = tuple(self._interface["shape"]) |
| 140 | self._interface["data"] = ( |
| 141 | self._interface["data"][0], |
| 142 | self._interface["data"][1], |
| 143 | ) |
| 144 | strides = self._interface.get("strides", None) |
| 145 | if strides is not None: |
| 146 | self._interface["strides"] = tuple(strides) |
| 147 | |
| 148 | @property |
| 149 | def __cuda_array_interface__(self) -> Optional[ArrayInf]: |
| 150 | return self.__array_interface__ |
| 151 | |
| 152 | @__cuda_array_interface__.setter |
| 153 | def __cuda_array_interface__(self, interface: ArrayInf) -> None: |
| 154 | self.__array_interface__ = interface |
| 155 | |
| 156 | @property |
| 157 | def shape(self) -> Tuple[int, ...]: |
| 158 | """Shape of the input array.""" |
| 159 | aif = self.__array_interface__ |
| 160 | assert aif is not None |
| 161 | return aif["shape"] |
| 162 | |
| 163 | @property |
| 164 | def size(self) -> np.signedinteger: |
| 165 | """Total size of the input array.""" |
| 166 | return np.prod(self.shape) |
| 167 | |
| 168 | arr = Array() |
| 169 | |
| 170 | # Cupy and numpy might run into issue when constructing an empty array from an array |
| 171 | # interface. we explicitly check for emptiness. |
| 172 | if "stream" in interface: |
| 173 | # CUDA stream is presented, this is a __cuda_array_interface__. |
| 174 | arr.__cuda_array_interface__ = interface |
| 175 | cp = import_cupy() |
| 176 | if arr.size == 0: |
| 177 | return cp.empty(shape=arr.shape, dtype=np.dtype(interface["typestr"])) |
| 178 | out = cp.array(arr, copy=not zero_copy) |
| 179 | else: |
| 180 | arr.__array_interface__ = interface |
no test coverage detected