If you're studying HPC, I would say that it's probably worth learning how to do it the hard way, especially if you will have other projects that require doing s

If you're studying HPC, I would say that it's probably worth learning how to do it the hard way, especially if you will have other projects that require doing so. However, if you want to use Python, I would consider JAX. It has a nice set of operations based on numpy and scipy and can compile to many backends (GPU, TPU, CPU, etc.) using XLA. The compiler is great at finding ways to optimize routines that you might not think of. Some of the manual parallelism functionality is still considered experimental, but I haven't seen that cause any issues or prevent functionality.