I think the bigger difficulty lies in how to store and process this stuff on a GPU/other accelerator.

If the OP cares about implementation details about how an API like PyTorch is made, I think the MiniTorch 'book' is a pretty good intro:

https://minitorch.github.io/