"UndefVarError: alloc_threaded_grad not defined" means that you weren't running the released version that is discussed in the blog post but the one of the much older slower ones. Any anyways, removing threaded calls and then benchmarking is clearly going to make it slower. With the updated ones, we seem to see 2x, 10x, and 22x, with increasing performance improvements based on the availability of threads and the existence of AVX512. For anyone who wants the details, just check that same gist:

https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...