JAX as benchmark

[pedronahum]

Hi all,

Wondering if you would be interested in benchmarking the performance of Swift vs JAX (similar to what you have done with Pytorch and Tensorflow).

Initial performance tests suggest JAX is much faster than Pytorch and TensorFlow, but still slower than Swift. Prelim results here: https://gist.github.com/pedronahum/61f9c96bcea6d6251a4406162d0fcdc1.

Best,

Pedro N.

[JaapWijnen]

Hi Pedro, definitely interested! I'll take a look after the holiday break!

[pedronahum]

Perfect! In the meantime, I will try to update the gist to include a TF version that leverages XLA. The message appears to be same, nonetheless, Swift looks to be faster than the AI/ML battle-tested XLA compiler in calculating gradients. But I let you do the official performance measurements :). Enjoy the holiday break.

[JaapWijnen]

Would you be interested in making a pull request?

[pedronahum]

Sure, give me a couple of days to move the code from colab and add a relevant folder with instructions.