| 130 | @numel_lookup.register_lazy("cupy") |
| 131 | @to_numpy_dispatch.register_lazy("cupy") |
| 132 | def register_cupy(): |
| 133 | import cupy |
| 134 | |
| 135 | concatenate_lookup.register(cupy.ndarray, cupy.concatenate) |
| 136 | take_lookup.register(cupy.ndarray, cupy.take) |
| 137 | tensordot_lookup.register(cupy.ndarray, cupy.tensordot) |
| 138 | percentile_lookup.register(cupy.ndarray, percentile) |
| 139 | numel_lookup.register(cupy.ndarray, _numel_arraylike) |
| 140 | nannumel_lookup.register(cupy.ndarray, _nannumel) |
| 141 | empty_lookup.register(cupy.ndarray, cupy.empty) |
| 142 | |
| 143 | @to_numpy_dispatch.register(cupy.ndarray) |
| 144 | def cupy_to_numpy(data, **kwargs): |
| 145 | return cupy.asnumpy(data, **kwargs) |
| 146 | |
| 147 | @to_cupy_dispatch.register(np.ndarray) |
| 148 | def numpy_to_cupy(data, **kwargs): |
| 149 | return cupy.asarray(data, **kwargs) |
| 150 | |
| 151 | @einsum_lookup.register(cupy.ndarray) |
| 152 | def _cupy_einsum(*args, **kwargs): |
| 153 | # NB: cupy does not accept `order` or `casting` kwargs - ignore |
| 154 | kwargs.pop("casting", None) |
| 155 | kwargs.pop("order", None) |
| 156 | return cupy.einsum(*args, **kwargs) |
| 157 | |
| 158 | |
| 159 | def sparse_take(array, idx, axis=0): |