Have you seen the gorgonia repos?

https://github.com/gorgonia/agogo

They basically went a similar route and had to implement all kinds of data types and GPU code to be able to train their alphago model.

In go, the biggest pitfall I'd imagine is memory management because you cannot influence the GC. So I'd build up everything with as many references from the start, and integrate it with the pperf tools to be able to benchmark it