Great basics tutorial: https://colinraffel.com/blog/you-don-t-know-jax.html
JAX – NumPy on the CPU, GPU, and TPU
81–90 of 147 posts
Re: JAX – NumPy on the CPU, GPU, and TPU
#82I’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_Gotcha...
Re: JAX – NumPy on the CPU, GPU, and TPU
#83Just going over the docs and there’s something I don’t understand, hoping someone here can explain:
What’s the point of having to explicitly call grad(jit(f))? Why doesn’t grad just call jit internally? Is there a usecase where you want the grad without jit?
Re: JAX – NumPy on the CPU, GPU, and TPU
#84Looking for some numerical problems to practice JAX, any suggestion or resource?
Re: JAX – NumPy on the CPU, GPU, and TPU
#85I’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 pur…
> 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.
Re: JAX – NumPy on the CPU, GPU, and TPU
#86Looking for some numerical problems to practice JAX, any suggestion or resource?
Although not terribly complicated, you can compute fractals like the Mandelbrot set or Julia sets. They look really cool and you can play with visualization.
Re: JAX – NumPy on the CPU, GPU, and TPU
#87> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code. Nice. When did they make this change? Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional. https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...
This is still the case afaik.
For vanilla "if", the condition must be known at compile time. For runtime, you have to use "cond", "where", or "select" (which may be analogous).
Re: JAX – NumPy on the CPU, GPU, and TPU
#88Earlier quoted context omitted.
I need the opposite - JAX to ONNX
I've been looking in to this for the java world. What's your use case? Deployment in to existing applications?
Yea exactly - Python for training, Java/.NET for inference at production. I looked at approaches like GRPC and things but my case is a bit more time-sensitive and the latency added by going over a network layer was too much.
For now I'm happy with Pytorch->ONNX and then running the ONNX model directly. But as I said, that means I can't easily train using JAX :-(
Re: JAX – NumPy on the CPU, GPU, and TPU
#89Earlier quoted context omitted.
the issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)
Jax JIT of scan is fairly good, so loops aren't as slow as you'd expect.
A key difference is that each iteration of scan is called by the host. Put differently, JAX can't fuse scan into a single GPU kernel, but launches a kernel for each iteration.
Depening on the workload this is no problem. If you have many cheap iterations, you will notice the overhead.
I am not sure if they are working on fusing scan and what's the current status.
Re: JAX – NumPy on the CPU, GPU, and TPU
#90I have been working on my DNN model using TensorFlow even though ML is not my main research. But it is a substantial part of my research, so I have to figure things out on my own, and I have done so over the past 3 years. However, I spend so much time on figuring out how any of TF methods works and debugging them. I never used JAX but I am not sure if this sort of grinding is normal when you use JAX as well (I always…
Is there a reason you are limited to JAX/tensorflow and can't use pytorch?
For a lot of people I know whose main job was not to write code, switching from tensorflow to pytorch was something that saved them ten to hundreds of hours in the long run, even accounting for the initial learning time.