have you looked at PyTorch ? can be seen as a GPU acceleratd numpy

see https://vwrideout.github.io/QR-Blog/index.html