Performance issue

Hello everyone,

I’m testing JAX framework speed and for that purpose, I wrote the Kalman filter. The gold standard is a Matlab version that performs an algorithm for 50k number of steps in ~0.1 sec. The TensorFlow 2.1 with tf.function gives ~2.0 sec. Unfortunately, JAX version hangs forever.

The JAX version of Kalman filter in Colab is here and TensorFlow 2.1 version is here. I would appreciate any help in resolving the issue or understanding what is wrong with the code.

Thanks, Artem

[skye]

I think you need to use block_until_ready in your colab microbenchmark. Something like:

%timeit -n10 kalman_filter_jit(y, var_y, var_f, len_f).block_until_ready()

This is because JAX dispatches operations asynchronously (usually a good thing!), so you’re only measuring the dispatch time here. See https://jax.readthedocs.io/en/latest/async_dispatch.html for more details.

Regarding the unrolled for-loop, this is an artifact of how JAX traces functions to be jitted. It’s a little nuanced, but basically, JAX only “sees” certain overloaded operations (e.g. jax.numpy functions), and Python doesn’t provide a good way to overload control flow, such as for-loops. The result is the loop runs as usual (JAX doesn’t “see” it at all) and JAX traces the operations inside each iteration, but can’t tell they’re all from the same loop body. That’s why we provide control flow operations like scan, so JAX can “see” the loop itself. This talk goes over JAX’s tracing in more detail if you’re curious.

[gnecula]

I am closing this, assuming that the issue was clarified. One lesson is that it would be helpful for the documentation to make it very clear that JAX is meant to optimize certain kinds of Python codes only.