Use XLA compilation with tf.function

dionhaefner / pyhpc-benchmarks

A suite of benchmarks for CPU and GPU performance of the most popular high-performance libraries for Python :rocket:

The Unlicense

310 stars 25 forks source link

Closed dionhaefner closed 4 years ago

dionhaefner commented 4 years ago

See #4

dionhaefner commented 4 years ago

Seems like this makes TF a lot more competitive on CPU! (and very close to JAX)

I'm surprised though that JAX still seems to be quite a bit faster on GPU 🤔

mdanatg commented 4 years ago

cc @cheshire Yep the GPU numbers are interesting...