I've made some effort to provide a set of similar high-quality implementations available in PyTorch: https://blog.millionintegrals.com/vel-pytorch-meets-baseline...

In my opinion PyTorch code is easier to understand and debug for newcomers. Code is definitely lacking in documentation, but whenever there was a tradeoff between clarity and modularity in the end I've chosen modularity. Ideally I would like others to be able to take bits and pieces and incorporate into their projects to speed up time to delivery of their ideas.