You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I was just comparing the performances of different ML libraries (basically to see if JAX is faster than PyTorch or Tensorflow - the consensus is that JAX is faster but I want numbers) and I noticed that the function generated by eqx.filter_jit was quite slow? It's easier to show in code. So, this is the very simple benchmark:
the pip installs to reproduce this
pip install equinox tqdm polars
The timeit util, not the main point of the question, but here for completeness
So this left me a bit confused, because I always thought that eqx.filter_jit was just a thin wrapper around jax.jit but that wouldn't explain the large difference. My tests were performed on a Macbook M1 on the CPU.
The text was updated successfully, but these errors were encountered:
Equinox does this so that runtime errors are correctly surfaced during the JIT'd call, and not at some later point (or possibly not at all if the program stops before then).
I tested your benchmark with this addition, using both jax.jit and eqx.filter_jit, and get comparable timings.
Hi there,
I was just comparing the performances of different ML libraries (basically to see if JAX is faster than PyTorch or Tensorflow - the consensus is that JAX is faster but I want numbers) and I noticed that the function generated by
eqx.filter_jit
was quite slow? It's easier to show in code. So, this is the very simple benchmark:the pip installs to reproduce this
pip install equinox tqdm polarsThe timeit util, not the main point of the question, but here for completeness
eqx.filter_jit
gives this result:whereas the
jax.jit(lin.__call__)
and the Flax version give these results respectively (which are besically equivalent):So this left me a bit confused, because I always thought that
eqx.filter_jit
was just a thin wrapper aroundjax.jit
but that wouldn't explain the large difference. My tests were performed on a Macbook M1 on the CPU.The text was updated successfully, but these errors were encountered: