3 ms·

Anyone from the Numba team care to comment?

• albertzeyer 7y ago
  While there is an overlap in how they work technically, I think the overlap is not too big. So if you ask about a Numba vs JAX comparison here, I'm not sure if such a comparison makes too much sense. JAX core is an extensible system for transforming numerical Python functions. This core is used to implement automatic differentiation, translation to TF XLA, etc. Numba does not have such a generic function transformation framework - it just supports a single transformation, that is from numerical Python function to machine code.
• cl3misch 7y ago
  A big difference is jax having a tracing compiler while numba does not. In general, jax is more suited towards vectorizable code (so using numpy functions broadcasted over axes) while numba is better for accelerating manual for loops. (I'm not from either jax or numba, but a keen jax user for non-ML research.)
  • hedgehog 7y ago
    Out of curiosity what are you using JAX for?
• emu 7y ago
  We (JAX) see JAX and Numba as mostly complementary; they have different strengths and they are focused on different things. We haven't tried combining them yet, but we think it would be fun to explore (https://github.com/google/jax/issues/1870 https://github.com/google/jax/issues/1870). For example, you could use Numba to hand write a numerical kernel that then participates in a machine learning model that uses JAX automatic differentiation.