5 ms·

I’ve used Numpy with Numba primarily (on the CPU) and it had been a game changer for my data science workloads. Naturally, I pay close attention to Jax and giv

I’ve used Numpy with Numba primarily (on the CPU) and it had been a game changer for my data science workloads. Naturally, I pay close attention to Jax and give it a look every now and then. So, I’ll focus my observations below on Jax’s Numpy API support. At a glance, Jax code looks like regular Python, but it’s a very different style of programming. Two big differences I’ve found are: - All Jax functions must be pure. You can’t pass references. - ndarrays cannot be created with dynamic shapes. You have to hardcode the shape tuples. One possible workaround can be to create a buffer much bigger than you need and return that along with actual shape. Then there are many small things that are very well documented[1] by the Jax team. Overall, if you are training ML models, the trouble might be worth it (Autograd). But for accelerating Numpy alone, it is no Numba replacement - which will happily work in the above mentioned use-cases. [1] https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...

cl3misch 3y ago
> You have to hardcode the shape tuples. This is not true. Rather, all shapes have to be known at compile time. That means that output shapes must not depend on input values, but may depend on input shapes -- also explicitely. Furthermore, there are two useful additions: 1. You can use vanilla numpy for compile-time computations. An example would be computing an array of indices for some moving-window filter, depending on the input shape and a stride parameter. 2. You can mark function arguments as "static". Then their values may change shapes of the output, but accordingly the function is compiled for each value of those arguments.
- sheepscreek 3y ago
  You’re right, I apologize for wording it incorrectly. It might have been a restriction with the # of dimensions. Either way, it wasn’t cut out for my use case.

cl3misch 3y ago

> You have to hardcode the shape tuples. This is not true. Rather, all shapes have to be known at compile time. That means that output shapes must not depend on input values, but may depend on input shapes -- also explicitely. Furthermore, there are two useful additions: 1. You can use vanilla numpy for compile-time computations. An example would be computing an array of indices for some moving-window filter, depending on the input shape and a stride parameter. 2. You can mark function arguments as "static". Then their values may change shapes of the output, but accordingly the function is compiled for each value of those arguments.

sheepscreek 3y ago
You’re right, I apologize for wording it incorrectly. It might have been a restriction with the # of dimensions. Either way, it wasn’t cut out for my use case.