After my last post on Macs and CUDA I was sent a link to CuPy which is a library that is supported by NVIDIA that allows to easily run CUDA code in Python using NumPy arrays as input.

CuPy's interface is highly compatible with NumPy; in most cases it can be used as a drop-in replacement. All you need to do is just replace numpy with cupy in your Python code. It supports various methods, indexing, data types, broadcasting and more.

To install

pip install cupy

Note

The latest version of cuDNN and NCCL libraries are included in binary packages (wheels). For the source package, you will need to install cuDNN/NCCL before installing CuPy, if you want to use it.

Or you can install versions specific to the particular CUDA environment. Full details are on GitHub https://github.com/cupy/cupy.