Why is my simple Jax program so slow?

[NeilGirdhar]

Simple functions that just apply matrices, nonlinear functions, etc. seem to be extremely slow. I'm not really sure if it's because Jax is not able to fuse kernels? What should I do to make this program faster?

Attached the Perfetto trace: perfetto_trace.json.gz

Functions like to_exp look simple to me, so I'm surprised that they're so slow.

Any help is greatly appreciated!