Machine Learning with JAX - From Zero to Hero | Tutorial #1

Key Takeaways

What's cracking, guys? With this video, I'm kicking off a series of video tutorials about JAX. So, what is JAX exactly? Let's see the official docs here. Uh basically, JAX is a autograd and XLA, which stands for accelerated linear algebra compiler, which is basically a compiler developed by Google that co-evolved with the development of TPU units. We're going to see a lot of XLA later, so I thought just kind of explaining the terminology. So, they are basically brought together for high-performance numerical computing and machine learning research. So, in a nutshell, we can treat JAX as a like a machine learning library and you need to be cautious not to compare JAX directly to PyTorch or to TensorFlow because JAX is kind of like a mid-level like library and people have been building on top of it. So, we have as for the high-level deep learning libraries, we have the two most popular ones are Haiku coming from from DeepMind and we have a Flax coming from the Google research team. So, basically aside from that, DeepMind team has been developing for every single like domain particular domain such as graph ML or RL, they've been developing like a purposeful dedicated library built on top of JAX. So, for example, for graph ML, they've built up a thing called giraffe or JAX for graphs. That would be like a acronym pretty much. Having mentioned Flax, I think it's worth noting that on Hugging Face Hub, you have more than 5,000 models at the at the time of this recording that you can use, download, and play with them and I think most of them were written in Flax and that's the reason why I'll be covering both Haiku and Flax in some of the future videos. That out of the way, let me tell you what I'm going to cover in this series of videos. So, in the first two videos, I'm going to explain the nitty-gritty details of JAX and then we're going to do we're going to code a neural networks from scratch in both pure JAX as well as Flax as well as Haiku. Also, those of you who know me, you know that I'm a big fan of a top-down approach when teaching and basically here I'm going to go a bit differently because the the topic is kind of complex and the first two videos are going to go so the first one is going to go from mid-level of abstraction down to lower lower levels of abstraction and we're going to understand the nitty-gritty details of JAX in this video. The second one will be mid-level to high levels of abstraction and as I said, the last three videos they will be hands-on and those are going to be on a high level of abstraction. We won't be digging deeper into the primitives, we're just going to build from the components we learned in the first two videos, we're going to build neural networks from from scratch. So, aside from that, as you can see here, I have a code skeleton already here and the there is a reason behind that. First, if I was writing if I was coding everything from scratch, that would take a this video would be 2 hours long. Secondly, this is way more natural because you usually never ever build build stuff from your head from scratch. You take you either take some snippets from the documentation, from Stack Overflow, or from some somebody else's repo. So, it's very I want to stress the understanding here. It's important that we understand the code so that we can modify it and I think it's less important for you to understand how to code this up from your head because nobody does that. I mean, sometimes I even like forget how to write down if name equals main in the main like file of a Python program. So, you just search for stuff if you don't understand how to write something down. Okay, let's dig down into the code. First off, I'm in Google Colab here. I enabled the GPU accelerator in the background and this Jupiter file, which I'm going to share with you after this video, I'm going to open source it on my GitHub, is structured into two like sections. The first one is warming up with JAX. Basically, we're on the mid-level of abstraction pyramid, let's call it that way, and then we're going deeper understanding the nitty-gritty details behind many of JAX primitive transform functions such as JIT, grad, etc. We're going to see those in a second. Okay, let's let's let's start. So, first things first, I want you to understand that JAX is JAX's API is very similar to NumPy's and second thing I want you to take out from this section is that JAX code is accelerator agnostic and we're going to see what that means exactly. Okay, let's let's dig into the actual code. So, as I said, for the most part, the syntax is the same as NumPy's and you can see here we're importing from JAX. There is this module called NumPy and the convention is just to import it as JNP. They also have a SciPy API so JAX SciPy, but we're not going to use that one in this video. So, these functions that are called transform functions are a vital component of JAX. So, those are grad, JIT, VMAP, and PMAP and you'll be seeing these a lot throughout this this notebook and in general using JAX. So, aside from that, JAX is structured so the API is structured like an onion as usually software like libraries are. So, we have the the high-level API, which is the NumPy, then we have this thing called called LEX and finally, we have I guess XLA, which is I think in C++ right? Okay. So, because I'm a I'm a etymology nerd, I just thought writing this down. Basically, LEX is just an anagram for XLA, which is a compiler and I'm not completely sure how people came up with the name JAX. So, if anybody knows, feel free to type it down in the comment section. Let's see a couple of facts about JAX. So, the first thing I already mentioned, the syntax is remarkably similar to NumPy's. If I run this, we can see we have we have defined 1,000 points on the x-axis equally spaced from 0 to 10 and I just plot the function two sine cosine and we just plot it we we get the chart here. You can see that JAX syntax is completely similar like it's actually same. If you just switch MP for JMP, you get linspace, sine, cosine, plot. You can like as you can see here, the arrays can be directly inserted into the matplotlib's plot function. So, if I run this one, we get the same results. Okay? So, that's it. So, that's the first fact. Second fact, so that's something you you need to get comfortable with and that's the functional programming paradigm, especially if you're only familiar with the object-oriented programming paradigm. This is going to be a bit rough, but like I I I promise you there are some awesome benefits that functional programming brings with it. So, JAX arrays are immutable. So, that's one of the three like that's one of the common properties you see in functional programming in in functional programs. So, what that means is the following. Um If we take a like a NumPy array here, so numbers from 0 to 9, we print them out and I then modify the the the the array at index 0 and I add some random number like 23 and if I run this, everything is fine. We have the original array and we have the modified array and as you can see, X is modified in place, which means the array is by definition mutable. On the other hand, if we take this the same program and run it here, so I create this time JAX array instead of NumPy array and I try to modify the array in place, this is what's going to happen. So, we're going to get an exception here as you can see. Basically, JAX is complaining that we says here object does not support item assignment. So, you cannot basically modify the array in place. So, this is the solution. It's kind of rough, but yeah, you have to do this at syntax with index and then set the value and this is going to work. Finally, you can see that we are not like modifying X in place. We basically allocate additional object in the memory called Y and now we have both the original array as well as the novel array. So, you may think here well, this is kind of suboptimal, right? Well, the trick is and we're going to see that a bit later. Like JAX uses something called JIT, which uses something called XLA, which basically takes care of this stuff behind the curtains. So, you don't have to worry about this. So, if it notices that you're not using the original array it's going to just do the same place and avoid allocating another memory object. So, yeah, that's taken care of. So, that's the the the second thing I wanted you to understand. So, the functional programming paradigm is something that JAX relies upon. Finally, we have the fact that JAX handles random numbers completely differently compared to NumPy and there is a very very good reason for this. We're going to see that in a in a moment, but the the the key point here is you have to create this basically key, which is a fancy name for for like a state. So, whereas NumPy's pseudo-random number generators are stateful, JAX is not stateful because it's as I said, it's that's a consequence of functional programming paradigm again. So, key is just a synonym for for a state in JAX and what you have to do when you're generating random numbers, you have to pass the state explicitly and not implicitly such like in NumPy. So, here we created 10 like random numbers. We sampled them from a normal distribution and we can see here that we get 10 random numbers. The thing I want you to notice here is that the type of this X array is of device array. So, that basically that's the that's the trick and that's why JAX code is accelerator agnostic. Basically, JAX automatically puts this this array onto the accelerator. So, in this case, I'm using here my accelerator is set up to to GPU. If I were to set it to TPU, the code would be set directly to the TPU unit. And that's very cool. Uh we're going to see a consequence of that a bit later. Basically, you don't have to do the laborious to device from PyTorch syntax, etc. So, this story uh naturally leads me to effect number four, and that's that JAX is accelerator agnostic. Uh same code runs everywhere. You don't have not have to uh write particular syntax for a particular accelerators. That's very cool. Um especially in the future where we we see more and more companies being uh like building custom AI accelerators and chips, not just TPUs and GPUs, but like companies such as Graphcore, uh Cerebras, etc. Let's see it in practice. So, here we define a JAX uh array. Uh again, we we see the syntax where we have the key. We have to pass explicitly the key or the state of the of the pseudo random number generator. And uh as you can always there is some difference in the API design here. We have D type here. We're using for NumPy, we're using S type. So, there are some minor differences, but like in general, syntaxes are very very similar between uh JAX's NumPy and NumPy API. Um I mentioned that we don't have to use to device anymore. So, that's in PyTorch. Uh so, this array here is directly pushed uh to the GPU unit. Whereas, this one is on on the CPU. Um if we do some time profiling here, we can see that if let let me run this thing. Um basically, uh this this first line line number 10, uh we we do uh like a basically matrix multiply directly on the GPU, and that's very fast. We can see that we're going to see a result in a second. Uh the second line here, uh because we have a NumPy uh dot and because we have NumPy arrays which are on CPU, uh it's going to be way slower. And NumPy only works The the reason why JAX exists is because NumPy only works with CPUs, and we do need to leverage existing accelerators. Uh the third line here uh does something like a like a modified uh version of the two of the first two lines. Basically, we're using uh JAX's dot product, and we are using um NumPy arrays which will cause this line to push the devices to the GPU, and and then do the the dot product. And because of that overhead, we're going to have a slower result hopefully for this third line. Uh and let's see the numbers here. So, we have 27 milliseconds for the first line. We have 437 milliseconds for the second line because as I said, NumPy is very slow. It's running on CPU. Finally, we have uh 97 second milliseconds for for this line here because we have the overhead of sending the data from the CPU from the host to the GPU unit. Um and the final number 26 is the same as this one, and that's this one here. So, we basically what we do, we we can explicitly using this device uh put function, we can push the the the the array to the GPU, and then uh now we have equivalent lines between this one and the line number 10. Uh and that's it. That's why we have the same results here. Um a couple of notes here. Uh first, I'm using GPU uh as a synonym for AI accelerator. In reality, as I said, uh JAX is uh accelerator agnostic, which means um basically, depending on what I set here in the runtime uh like setting here, uh we we are going to run this either on TPU, GPU, or whatever they have supported. Uh second note is this block until ready. You can notice that every time I'm using JAX's uh dot function, I'm using this block until ready. The reason is uh JAX is using something called asynchronous dispatch system in the background, which means if I were to run this thing here, so let me just copy paste this. So, if I were to copy paste this here, uh if I were to run this, and uh maybe like let me assign that to a variable, and now I do some coding here, blah blah blah. Uh the thing is this line is going to immediately return, and this code here is going to start executing because this line actually just delegates the task to the accelerator, and we're not blocked. We're not waiting for this line to execute, which means the the the variable te- temporary ver- This temporary variable is still not uh filled in. Let's let's put it that way. This is a fairly neat feature that makes things even more even faster, but you need to be cognizant of it, especially when you're doing this time profiling. You want to measure the actual computation, and not not the time it it's it's ne- that's necessary to dispatch just delegate the task to the accelerator. Just be aware of that. Nice. Those were some basics. Uh now we're going to cover the transform functions. Uh but before that, let me do a quick overview. So, JAX is AI accelerator agnostic. Uh JAX handles random numbers a bit differently. We're going to see why that is important a bit later. Basically, reproducibility in the environment where you have multiple uh accelerators, where you're doing parallel programming, uh is um like possible with JAX, whereas NumPy has some problems there because it was designed mainly for CPU programs. Uh then, as I said, uh functional programming paradigm is something you need to get comfortable with. And finally, um basically, the syntax is very similar to NumPy's. That's a high-level overview of what JAX is. And now we're going to get deeper uh into these uh transform functions. So, as I said, JIT compiles your functions using this XLA, and basically, uh doing that, it caches some functions, and makes it very very fast to execute our programs. Uh I'm going to just run this simple visualization function here, and I'm going to show you a simple example where how we can JIT a function and make it faster. So, here I just defined this function. Uh it's a Celiu function. It's just an activation function. The details here are not that important. That's why I have this visualize uh like function here. I'm just going to run it. Uh we're going to see how it looks like. You can see here uh like up until zero, and that's like I mean it's defined here. When we are when X is greater than zero, we just have X, which means we have a linear function here. And uh like when we are when the number X is smaller than zero, we have some combination of exponentials here with some coefficients, etc. Uh but that's not important. The important part here is we can transform Celiu. So, we pass the function inside of this JIT transform transform function, and we get a function back. And this function is compiled. Actually, it's not exactly compiled. We first need to trace it by calling it once, but like for the sake of argument, we just have a like a super fast function right now here. Uh and let's benchmark it. So, let me uh allocate a vector of million data points. So, we have million random random uh like numbers here, and we pass them through the Celiu function, and we basically um do a time profiling of both the normal version as well as the JIT version. Let me let me run it again, and let's see the numbers. So, uh the the the thing here is this function is not overly complex. So, the results we get with JIT, so we we will not get a huge uh performance boost here. But imagine once you start training neural networks, that's where uh JIT starts shining because then the the optimizer the compiler has a uh bigger freedom to do various various uh like optimization tricks such as fusing the operators, such as um avoiding allocating certain temporary structures, etc. And things get really really fast, like orders of magnitude fast faster. Here we can see that uh basically, we have uh 1.96 milliseconds using the normal version, and then we have only 121 uh microseconds using the JITted version. So, even in this simple example, we get performance benefits. So, for now, um I'm going to leave it at here. I'm going to stay like a on a high level here. Uh we don't have to understand currently how JIT works, but like you you just need to know when you see JIT, uh JIT makes functions run fast. Let me try to beat this guy. Awesome. Okay. Let's let's uh go ahead here. Um second transform function, very important, is grad. Uh basically, it does the uh like the magic of automatic differentiation for you. So, the same thing as dot backward in in PyTorch if you're coming from the PyTorch world. Um and uh like a short uh like a note here, uh differentiation can be manual. So, basically, you have a function, you manually calculate on the paper uh how the uh like how the derivatives look like, and then you can code those that knowledge into a function, and that's a manual differentiation. Symbolic one is very similar. It's just an automatic way to to automate the manual differentiation, whereby the program is using those rules such as product rule, etc., to build up uh like derivatives of a function. And numeric function are things like finite derivatives, where you use basically numeric uh methods to compute the derivatives. And finally, the automatic one is the one we all love, uh which is used in every single deep learning framework we know of. Um okay. Let's let's see how grad works. So, let's define a simple uh function here. So, it's just a sum of, as you can see here, logistic functions. So, this here inside of this is a logistic function or sigmoid. And um I just input I just created a like array of uh three values here, 0 1 2. And uh I renamed the function to loss because uh this could be used as a as a loss function, just giving it some semantics. And finally, uh we can do grad by just uh wrapping loss into grad. Again, we're passing function into a function, which is something you see often in the functional in the functional programming paradigm. Uh that gives us the grad loss, so the the gradient of the loss function, which we can evaluate. So, this is contrast this to to to PyTorch or TensorFlow, where you just do backward. Here, you actually get a function back, and you can evaluate it at particular points. Uh and uh by default, I I said here the the the grad will take the derivative of the first parameter, but here we only have one parameter, so that doesn't matter. And let me kind of run this, and see what we got. So, we have some numbers, which mean nothing to us because the function is fairly complicated. If I were to change this to this sum of squares function, so if I do something like this, we basically can get interpretable results. So, let's do a manual derivative of this function. We're going to get So, that's X1 squared plus X2 squared plus X3 squared. Remember, we are passing three numbers here. So, if we do derivative of that, that's going to be X1 * 2 + X2 * 2 + X3 * * 2. And if I were to run this, we're going to get 0 2 4, which makes sense because we have 0 1 2 Let me print this out. So, print X. So, we have 0 1 2 and we get 0 2 4 because of this. Grad basically does because all of these are bundled inside of X. Grad basically does derivative of the function with respect to X1, which is 2 X1 and then the same for the other variables. Originally, remember we have something like this. We have this thing plus X2 squared plus X3 squared. That's the original function. This here is the derivative, manual derivative. Okay. That's a simple example. Let's now go and make sure that this thing actually does what we expect it to do. We can do that using this numeric differentiation method, the finite differences method. If I were to run this, we're going to get some results. There there are some small errors here, which are natural consequence of the fact that we're storing numbers using a finite number of bits. So, how this works is very simple. We We We take So, remember, X is the this array here, 0 1 2. Length of X does will return us three. We get the I is just the identity matrix. And so, we are basically iterating and taking one-hot vectors here. And basically, this is here is the This is the formula of the actual derivative. So, you you nudge the input vector along a certain dimension a little bit in the positive direction, in the negative one, and then you divide that by the two epsilon because yeah, that's the intensity of the nudge. And if you don't understand this, the best way to play with Colab is to just insert a new code line here. We can see what this exactly does. So, let me just paste this here. So, we we have this thing here. So, we have, as you can see here, identity matrix. And so, this for loop is going to take one vector at a time and pass it into this function here. And so, you can see that basically we're going to evaluate F at the following up. So, at least in the first iteration, we're going to evaluate it at zero, and then we're going to add the small epsilon here, so epsilon, and then we're going to have one two. And we're going to evaluate the function here. We are going to subtract from that X minus epsilon. So, sorry. So, that's zero minus epsilon. This thing will be constant for the first iteration, and then we divide this by two epsilon, and you can see that this is the definition of derivative itself. And we do that for alongside every dimension. That means in the second loop, we'll have epsilon here instead of here. And so on and so forth. So, minus epsilon here, and yeah, you get the point. So, now we are certain that Grad works as expected. Now, let's see some some fun examples. We're going to define a simple second-order polynomial function here, X squared plus X plus four. I'm going to run this to visualize it. It's always nice to visualize stuff. It's at least to me, it makes stuff a bit more like a It gives them this gut feeling, and it's easier to understand what's going on. We're going to do higher-order derivative here just to show you how powerful Grad function actually is. So, you can do Grad of F, and then you can do Grad of of the Grad of F, and then you can do Grad of and you can do that N times, whatever the number like whatever the whatever N is. Or you can just do it like this. You can do it like Grad of Grad of Grad of F. This will also work. This is just shorter, so that's why I did it like this. If we were to manually do a derivative of this polynomial, we'll get the first derivative will give us 2X + 1. The second derivative So, derivative of this one will just leave us with two. And finally, derivative of a constant is zero. That's why we have the third derivative equals equals zero. So, if we print all of the values here, we expect to see because X is one, we expect to see three, two, and zero, right? So, here it is, three, two, and zero. The first value is just the value of the function at the input X, which is one, which is, as you can see here, approximately six. Exactly six. Okay. So, so far, so good. The cool thing about this is we are very close to the math. So, you can basically, as you see formula in some paper, you can implement it much easier in JAX compared to PyTorch or TensorFlow. It's very powerful. We're going to see how powerful this whole auto diff package of JAX is a bit later. And now I'm going to do a simple modification here. Let's assume we have two inputs. So, let's assume we have Y here. So, we have plus Y squared. And I'm just going to modify this lambda function. I'm going to skip the evaluation because the evaluation will fail now. And now I'm going to do the following. If we just do grads of the function, grad will, by default, do the derivative with respect to the first parameter. That means we're going to get this thing here exactly. So, we we should expect exactly the same numbers except for the fact that F will be different when we evaluate it at the at X. So, this obviously will not work until I enter some numbers. So, this is going to be like let me let me say it's going to be one, and then I'm going to pass modify this here. So, one X Y X Y, and finally X Y here. That should work. Let's see if it works or not. Yeah, we get the numbers as expected. So, we get the same results here, and we get seven because once evaluated at Y equals one, this will be 6 + 1, that's seven. So, how can we do derivative with respect to Y? The only thing we need to do here is arg nums equals one instead. I think this is the syntax. Let me try it out. Let me just ignore this for for for the time being. I'm just going to delete this, and let me print this out. So, we get the number two here, which is correct result. If we do a differentiation of this polynomial with respect to Y, we expect to get two times Y. Since Y is equal one, we get two here. So, the thing that confuses me here is I'm not sure why we're getting device array here, whereas previously we had pure floats on the CPU. I'm not sure about that. Let me let me check. Let me check if I can do it like this. Yep. For some reason, once I add like comma here, it's returning me device array. Whatever. Obviously, this is now DF DY, but like yeah, you get the point. Now, let me show you how powerful the auto diff engine is. So, aside from grad, which can do derivatives with respect for the scalar output functions, Jacobians, so these JAX forward and JAX rev functions can find Jacobians. So, Jacobians basically can evaluate derivatives even for vector-valued functions. That's the only difference in That's pretty much the the only difference. So, here let me let me take another function. Basically, we have a simple paraboloid. If we Google it here, we can see how it looks like. It's like a cup to the surface in a 3D space. And basically, if we manually calculate the derivatives, so DF DX will be 2X, DF DY will be 2Y. So, we expect Jacobian to look something like this. If we continue on and and and figure out the second-order derivatives, we can get the Hessian numbers. So, we can get the second-order derivative with respect to X is two, as you can see here. For Y, it's also going to be two. And if we do DX and then DY, because this does not depend on Y anymore, we're going to get zeros. And finally, how you form the Hessian matrix, and you don't need to worry about what Hessian is. It's basically just It's just a collection. It's It's a matrix of, as you can see here, various derivatives of the multi- variate function. You can simply define it using the JAX rev and JAX forward function. We additionally jit it. You can see how this nicely composes, and we can get the results. So, the reason they're using both forward and rev is because it's just optimization like detail because one of these, the rev one, works with wide matrices, whereas the JAX forward works with with tall matrices. I'm going to link a video down in the description, which nicely explains why this is. Anyways, we get the numbers, so that's two two for the Jacobian, as we expected here. So, if we plug in the numbers, so for the input 1 1, we can see that this evaluates to 2 2. And finally, for the Hessian, we get 2 0 0 2. So, anyways, um I just wanted to show you that this is possible, and if you want to dig into more details, their documentation is very very nice, and yeah. Let's continue with the final example for grad. I just took this edge case function, so the the value of X, and let's see how how how JAX handles it. Uh basically, this is how the function looks like. I just visualized it here. And um I printed the the values at -1 and 1, which would be uh 1 and 1. And you can see the the the results here. And finally, I printed the I I found the gradient of this function, and you can see it's not differentiable at this point at zero. It's that the gradient is undefined. Uh but like if we were to evaluate the function at uh -1, we get Um so, as you can see here, the derivative will be -1 for all the numbers here, 1 for all the numbers here. The interesting point is actually zero. Because as you can see, this one returns the obvious result, -1. So, this one here is kind of interesting. If we were to delete this part and just kind of uh add a small number to to to to zero, we'll get uh 1 as the output. If we were to do the same thing just on the negative side, we're going to get -1. So, basically, what they've done is they defined for zero, it's going to be the value is going to be uh defined is going to be 1. That's just a convention, and that's how JAX deals with it. Uh we didn't get any exception, so yeah, I guess I I I guess that's good at at least in some settings. Finally, let's jump to Vmap. Um basically, when you see Vmap, you need to think about it. The the the main value prop of Vmap is the following. Write your functions as if you were dealing with a single data point. And this is going to become uh more clear as we go through these examples. Let's say we have a matrix W, which is basically um weights of You can treat it as a weights of a linear uh neural network layer. We input some state. Uh we we set the shape to 150 100. And we now uh create this batched X. Uh you can treat this as maybe 10 images, a batch of 10 flattened images. So, it's 100, that means we had like a 10 * 10 pixel image. It's 100 when you once you flatten it out. So now, uh let if we apply the matrix, we do the dot product between W and X. Uh this is basically uh simulating what the linear layer is doing in the background uh when it's processing the input data. Now, the trick here is uh this only uh transforms the um the single a single image. It cannot handle the batch. Otherwise, it'll it'll crash because we're trying to multiply 150 100, uh which is W, and we're trying to multiply that with uh 10 uh 100, that's going to fail. So, this thing does not work for the batch. So, okay, let me let me run it. So, how would we go about uh making this work for for a batch of images? So, that's because that's what we care about. And because that's way more efficient, as you as you may know. Uh especially on GPUs and accelerators. Uh basically, the the most naive approach would be to iterate through the images and then call this this this function that knows how to do do to handle single data points. And then we just stack the results, and basically, we we can see that this will work. And the the thing is, it's very very slow because you you you you want to avoid doing for loops. You want to vectorize your your your functions, and that's why this runs 5.54 milliseconds. Now, let's see a bit more uh a better approach of doing this. So, this would be a better approach, but as you can see, uh we had to completely rewrite uh the function that we had above. So, we had to swap as you can see here, we had to swap W with X. Now, we have batch of X here, and we have W here. And we additionally had to do a transpose. So now, this will the shapes will match, and everything will work fine. So, we have 10 100, and we try to multiply that with transpose version of W, which is 100 150, and that gives us, as you can see, 10 150, which is uh what we expected, right? So, everything is fine. Now, the problem with this, and it's additionally jitted, so it's going to be really fast. Now, the problem with this is, as you can see, uh we have to write, depending on whether we handle singular cases or or or a batch of data, we have to write completely different functions. And that's not desirable. Although you can see this is way faster, so it's only 103 microseconds, whereas this is 5.54 milliseconds. Now, this is the result that JAX offers us using something called Vmap. You just take the uh function that handles a single data point. You wrap it into this You transform it using this Vmap, and you can now pass the batch of your data without any other modifications. So, that's that's very cool. Um if I run this, let's see how how how fast this thing is. So, it's around 100 microseconds. I think the last time I I ran this, this was actually faster than this one, but yeah, I guess there are some some some variance inside of there. Okay, so I went ahead and reran this thing again, and now we can see that we have 180 microseconds for this function, whereas we have 143 microseconds for this one. So, obviously, there is some variance in here, but the point is, here you have a much simpler way to write these batched uh functions, and it's as as efficient as by like doing this laborious work. So, uh basically, what what Vmap does uh in the background is it takes the the for loops and packs them into this LAX uh API mid-level uh layer of the of the API. Uh and yeah, so now, let me let me go ahead and um modify this example a little bit. So, we are going to instead of passing just X, let's try and pass both the W as well as the X, because that's more um I guess in in line with the functional uh programming paradigm. So, let me rerun this cell. And now, let's modify this one to to accept like both the W as well as the uh as the batch of data. So, uh we'll need to add this thing here again. So, if I were to run this uh right now, uh it will crash, and we'll see why. So, it basically says um it basically says that um W We're trying to to to to tell to Vmap that W has a batch dimension, where whereas it does not have. It's It's just a It's just a matrix uh that that's supposed to represent the linear layer. So, what we need to do is add the in axis argument here. So, in axis like this, and then we say none because W does not have a batch dimension. And we need to specify the batch dimension of the second input, and that's zero. So now, this should work. Let me let me try and rerun it again. So, fingers crossed. Um And yeah, it did work. So, basically, that's it. Now, you saw in a bit more detail how Vmap works. And this is the first part of this video. So, we basically saw the basics of JAX. We saw how to use the the main transform functions such as jit, such as grad, and finally Vmap. And now, let's dig uh even deeper and understand the intricacies of how how jit works, because that will help you enormously uh debugging these uh JAX programs. So, let's let's continue here. So, uh we basically have, as I said, JAX has this onion-like uh API layer uh API structure, and that's that's I guess pretty much uh always the case. But still, uh anyways, we have NumPy as the as the highest level, then we have LAX as the mid-level, and finally the XLA. Uh so, LAX API is stricter and more powerful, and it's a simple Python wrapper around XLA. So, let's see what it means when I say stricter. So, if we were to add uh in the NumPy uh level of the API, uh this thing can be tolerated. So, 1 + 1 uh like 1 as the integer plus uh 1 as the float uh will work. But once we get to the mid-level, uh here we need to be explicit about the types. Uh otherwise, we'll we'll have we'll we'll get an error. So, if I run this, we'll see that uh this is printed, so 1 + 1 is okay. But here, add requires arguments to have the same d types, got int32 and float32. The reason they've done this done it this way is to be uh more error error robust. Uh finally, uh we have uh the fact that the LAX layer is obviously more powerful, although as a trade-off, it's it's less uh user-friendly, which is kind of obvious. As an example here, we have X and Y. We have uh basically, uh we want to do a convolution a 1D convolution between the signals X and Y. And the NumPy uh API, this would uh be uh done like this. We just call the convolve function, and we get the result. On the other hand, once you get into the LAX land, uh you have to use this conv_general_dilated, which is way more powerful. As you can see, it has much more options. You can specify the window strides, the padding, but you have to be again much more strict here in order to get this to work. You have to be explicit about the the types, etc. And if I run this, uh we can see the results here as well as here will be the same. Uh like this this assert should make sure that that uh we get the same results. And as you can see, it works uh perfectly. Now, uh final remark here is, you can see here that uh this result re re uh returns uh batched results, so that's why we have to to index 0 0 to get the actual result. Uh as I said, uh LAX API is just a thin wrapper around thin. I mean, it's a wrapper around uh XLA. So, you can actually find the the XLA function that's going to be uh eventually called here, and that's this uh conv_general_padding. Uh it's in C++, you can see the arguments, etc. So, if you ever need to to to dig a bit deeper and optimize something uh really really really hard, uh then you you TensorFlow uh documentation here got got you covered. So, let's get back here and and continue. Uh that was the uh short uh mention of the API. Now, let's finally understand how jit works. So, uh again, uh I won't get into more more details here. The whole point of this is to show that uh JIT functions are are faster compared to uh non-JIT versions of the function. Here, we just normalize the matrix X like a column-wise, and so both the mean we subtract the mean as well as we divide with the standard deviation of the columns, and we get normalized columns. And finally, uh the difference here is again not that big because this is a super simple function. The more complex the function is, the the bigger the differences here will be between the JIT version and between the normal version. So, here we have the JIT version, uh and again, we're using a block until ready because of the asynchronous dispatch, and we see the results. Okay. So, let's see uh in order to understand JIT, it it may be useful to understand when it fails. Which functions cannot So, which class of functions we cannot uh JIT uh will tell you uh will help you understand how JIT actually works behind the curtains. So, if I were to run this thing, uh it's going to crash. So, we are creating a simple vector of 10 random elements, and we we we pass that vector into this uh get negatives, and we call it. So, now because we don't have any JIT, uh we're going to get a a result back, and everything works as expected. But, if I were to call uh the JIT version of the function by by wrapping this into this JIT transform, if we run this, we're going to get a like error. And it said it says here array boolean indices must be concrete, got shaped array, blah blah blah. So, basically what happens here is that depending on the content of X, so depending on the values of X, uh we're going to get uh an output that's going to vary in its shape. And that's something that's not tolerable tolerable inside the JIT world. Let's slowly understand why this is the case. So, let's start with this function. We have a simple function F which just does uh like a dot product between X and Y and returns result, and in the meanwhile, it also prints some some intermediate variables here. So, I I prepare the variable X and Y, just random vectors and matrix matrix here, and uh we we call the the function here, and then we again call we have a second call here, and let's see what happens uh after after the second call. So, I'm going to run this. There's a couple of things happening here. First things first is the first time you run a JIT function, so that that's line 14, uh what JIT does in the background is something called tracing. So, it basically takes uh instead of uh inputting the actual values of X and Y, it creates this abstract uh like a tracer value, let's call it that way. So, it's basically like a placeholder variable that has a specified shape and specified data type. So, you can see here, uh basically once we print X, it outputs here instead of the value of X, it outputs traced shaped array. It's basically float 32 because we had random numbers uh sampled from Gaussian. Uh we have the shape here, three uh {comma} four, and we have four for Y. So, basically you can see these are actually passed into the function the first the first time JIT is called, and uh using that, uh JIT can understand how the shapes are morphing uh going through the function, and finally what the output shape of the function is. And you can see here the output is also float 32, and uh the shape is three. We finally get the results here. The second time you call the function, something something funny happens, and that has to do with the functional uh programming paradigm. Basically, because print functions are a side effect, because we are returning the result from this function uh by other means uh than through the actual output here, so we are printing, and that's a side effect. And because of that, uh JIT is just going to ignore all of those side effects, and the second time you call it, because it's going to call the compiled function, uh we we won't see any printing. And the whole point is this this this caching mechanism I just explained. So, that's why JIT functions are so fast. So, I I mentioned here anytime we get the same shapes and types, we just call the compiled function. So, that means if I were to now call uh whatever So, so basically for the whole class of inputs that have this shape uh and this data type, and no matter the value, we're always going to call the compiled function. Uh but, if I were to call this function again, but this time with a different shape, so let's do it something like this. So, X3 may be three five, and this will be five, and let me just map this like this, and I'm going to call the function again and print the results. So, X3 Y3. So, what will happen here is that this time will will again trigger the compilation because the shape changed here, and so JIT is smart enough to to to retrace it, and as you can see here, we have the tracing happened here again. Um hopefully, you get a better picture of how uh JIT now works, but we we're going to continue and understand this in even more depth. So, now uh we have the same function as above. We just omit the print functions, which are the side effects, and this is how you should write your your your functions when you want to use them with uh JAX uh transform functions such as JIT. So, no side effects. Uh if I were to print this, and we are going to have something called JAXPR here. So, that's JAX expression, and that's basically the let's call it like a flow model that JIT uh creates in the background when it does its tracing procedure. So, if I were to run this, let's see what we what we have here. So, you can see uh the what what happens. So, it it creates this this abstract grammar, and you you basically have C, so add one to A, so that's going to be uh and A is basically the first argument, so that's X. So, it's basically creating a placeholder for this thing here. Then it calls D uh B plus one, so that's Y plus one. This is this is going to be B. And then it calls the the general dot product, so that's this thing with C and D, as you can see here, C and D. And yeah, and it returns back the E, which is this result here. So, you can go into docs and understand uh in more detail every single part of this of this syntax of this grammar, but like uh now you understand a bit better how JIT, uh when it does the tracing, uh creates this type of of uh of a flow in the background, and this can be compiled using XLA. Okay? Let's see another example of a failure. Um this time, what we do is we pass this argument uh neg, and we condition upon it, and depending on on on the value, whether it's true or false, it's we're going to return minus X or or plus X. And uh remember, uh the first time we call JIT, we call a JIT function, uh it's going to input abstract shape and and like a data type. It it won't have a information about the value, and that's why this thing is going to going to fail. So, it says here uh blah blah blah, um abstract tracer value encountered where concrete value is expected. Uh the problem arose in the in the bool function. So, that's basically what happens. And in order to avoid this, uh what we can do is is use these static arguments. Uh by making an argument static, so we say here, "Hey, first argument, so that's neg, is going to be static." Uh what we do by by doing this is once the JIT does the tracing procedure, it will not use this abstract tracer object, it's going to use the actual the actual value. So, we are kind of lowering the level of abstraction here while doing the tracing, which is going to uh kind of constrain the the class of inputs where this compiled function can be can be called from the cache, uh but as in return, uh yeah, in return we get it to work, I guess. Okay, so let's let's let's call this thing. So, this should now work. And the the thing I want you to understand here again is that the first time we call it with true here, uh we do the tracing, and then uh the second time we call it with true, because nothing changed, uh basically again, we don't have the tracing procedure. Once we switch this to false, we again trigger the tracing. So, now we can call this function for any integer 32, uh and we'll have two cached functions, so that's uh like which run really fast. Let's continue analyzing the failures. Uh the third failure is here. If I were to run this, let's see what we get. Uh blah blah blah, shapes must be 1D sequences of concrete values of integer types. Okay. So, what happened here is let me try and print some values again. Print X. Let's print uh X shape. Let's print um this whole product thing, and I think that's going to work. So, let's let's try and print that. So, as you can see here, what happened is that X is of trace type, and then X shape is actually like a concrete value, and then this thing here is again a trace, so this traced object, and we're trying to pass that traced object into reshape, which expects a concrete value. That's why this thing is uh crashing again. Let's see what the solution is. Uh we can basically uh use NumPy uh prod instead of uh JAX uh product. So, this may be a little bit confusing. Uh basically, as you saw, you have two tools to make these functions work with JIT. One is to make certain arguments static, uh and sometimes you'll have to use NumPy functions instead of JAX functions. So, don't ask me why. Um I'm only a couple of steps ahead of you. I I recently started learning JAX myself. But like, let's try and run this. And this time it works because we have the M prod will return a concrete value and not the traced the traced uh the traced object. And that's it. Um if you wish, as I said, I'm going to share the Colab. I'm going to uh push it to my GitHub. So, go ahead and play with these failure cases yourself. To better better understand how JIT works. But, basically, just keep in mind that JIT passes in the these abstract objects that have the shape information and the data type information and no value. And that will help you save yourself a lot of from a lot of headache, I guess. Okay. Um This is something I need to cover because there are some gotchas, some um idiosyncrasies of JAX which you need to understand to be able so that we'll have easier time in the later videos building up neural networks, etc. Once you understand a couple of these gotchas, things are going to be way easier. So, let's start with gotcha number one. Pure function. So, JAX is designed to work only on pure functions. We saw some glimpse of this while we were using print functions a couple minutes ago. And here is an informal definition of a pure function. So, first of all, all the input data is passed through the function parameters and all the results are output through the function results. Okay? So, that's the first thing. Secondly, a pure function will always return the same result if invoked with the same inputs. So, it's in a way, it's like a huge memory, like a cache table, and depending on inputs, you just retrieve the outputs. So, that's an informal definition of a of a of a of a pure function and I think it's going to be good enough to understand the following cells. Okay. So, let's start with the cell number one. Example number one. So, here we are violating number one because all results are output through the function results. Nope, that's not the case. We are returning some results over the print function instead of through X through through this return statement here. So, let me let me call this function. Okay. On the first call again, we have the tracing. So, we we call the executing function print but the second call due to the fact that we are passing the same shape and same data type. That's the just a float 32 here. We just use the cache function and we return five. Finally, the third call because we change the shape and the type here. So, we basically have array now. It's going to basically call and trigger the the tracer again and I think that's pretty clear at this point of time considering the other examples I already covered. Example number two. So, you do not want to interfere with the global variables and I think this is probably a bad design decision anyways. Even if we ignore the functional programming paradigm, this is probably not a good idea to do because yeah, unexpected things can happen with your program. So, what's happening here is that as I said, we are violating both one and two because let's see let's see why. So, all the input is passed through the function parameters. Nope, this time we are passing the input through the global variable. And secondly, a pure function will always return the same results if invoked with the same inputs, which is also not the case because depending on the value of G, we are going to return different values here even if we keep X the same. So, um what will happen here once we call the version of function the first time, it's going to cache the value of G, which is zero. Then we're going to update it and then we're going to call it with five and it's going to use the cached it's going to call the cached version of the function, which which uh if you recall has G equals to zero somewhere somewhere in the in the JAX part and that's why we're going to get the wrong result. So, let me let me execute this cell. So, on the first call, we get four and that's fine. Okay? Because we we passed four, we had G equals zero, that's fine. Now, G is equal to 10 and we have five here. So, we expect 15 but we get five because the G version of zero is cached in the in the function. And finally, if I call with a different we we change instead of float, we pass the like JAX array. This time we trigger the execution and we get the the correct result again. So, in any case, this is a wrong idea, a bad idea to to to add these types of of impurities. Example number three. So, Haiku flex are basically built upon this idea. The whole idea is that it's fine to have a stateful like functionality inside of your pure function. So, what I mean by that is the following. So, we have X here and um we do what we do here, we we create this state dictionary and we basically just add depending on whether the integer I here is odd or even, we add X to even or odd keys. And so, we are preserving state obviously by doing this. So, this is stateful execution here. And then we just return this sum here. So, as you can see, we're not violating anything here. So, this is this is this is pure because we are only using X. We are only assigning X to the state and we are only outputting whatever came out from this function. So, we're not accessing some some global variable or whatnot. So, because of that, all of this will be fine and if I execute this, we are going to get a correct result. And that's 50. Okay. Finally, a fourth example. Iterators, since they are stateful, are a no-no. So, if you try and do So, we built array here and if we just do like this, if we don't have any iterators, it's going to work. But, on the other hand, if we use the iterator, this is going to to fail even though we semantically we would expect the same results. And you can see that's the case. We got 45 here. We got zero here even though we we should have gotten 45 as well. Here you can see this this lax primitive. So, that's one of those like mid-level API functions. It's called fori loop. It's basically a smart version of the for loop which can be later compiled using the XLA. But, the thing I want you to to understand here is just that you cannot use iterators because they are stateful and we are thus violating the purity constraint of of JAX. Okay, that was the gotcha number one. Make sure to write pure functions if you want to use them with JAX transforms and in general. Gotcha number two. In-place updates. We saw this already. So, you cannot modify the the arrays in place. You have to to use that at set syntax. So, let's see a simple example here. We have to once we create the JAX array, we have to use this at set syntax in order to get the the the output array. Let me run this. And you can see here the results are as expected. So, this is your NumPy syntax. So, row number one where we have zero indexing and then all columns, we set them to one and that's what we see here. So, I think I mentioned this. So, if this seems wasteful, basically, don't worry because XLA is smart enough to figure out that once you're if you're not using this input array, then it will not allocate like a special memory object for this output array. It's just going to reuse the input array and modify it in place even though it does not appear to to do so on this higher higher level perspective, right? Don't worry about the expressiveness. We can still do everything we can with with NumPy. So, here if I create like a simple matrix of ones, we can do whatever we want. We can add. We don't just have to assign values to to certain locations. We can also do arbitrary operations such as addition. And you can see here we can add to every second row. So, every second row and starting from the third column, we add seven to one. So, that's why we have these eights here as you can see. So, that's cool. That was the second gotcha. Let's let's let's continue on. The third one is out of bounds indexing. So, this is a direct consequence of the fact that JAX wants to make this code accelerator agnostic because it's very hard to communicate certain certain information from the accelerators. They had to create certain types of non-error behaviors which may come as a surprise if you're not already acquainted with this. So, when it comes to NumPy, if we allocate an array with 10 elements, so from zero to nine, and we try to index into 11th position, this will throw an exception because there is no 11 position. There is position zero through nine. So, let me try and run this and we're going to get exception caught blah blah blah. Basically, the this is out of bounds for axis zero with size 10. Okay. Let's see what JAX behavior here is. If we were to try to assign at position 11, so that's the same example as here. If we try to add 23, it's just going to ignore ignore this this operation. So, that may be surprising because it will not throw any error. So, let's see what the result will actually be. So, we have no change whatsoever here. Finally, um if we try to return to retrieve the element at 11th position, which does not exist, it's going to clamp 11 to the last like index, which is nine. And so, we are going to retrieve nine, which is super confusing as well. And this may be a cause of many many really bad bugs. So, be aware of this behavior. As I said, the reason this exists is because basically, JAX tries to abstract the accelerator like information from you. But, as a consequence, you get this. And I guess if you're familiar with NaN behavior, this is kind of similar because NaNs are also in a sense not acceptable, although we we just we don't have the this we we we also have non-error behavior, whereas we just get NaNs back instead of like the system throwing some exception. So, I mentioned that here. Similar to how invalid invalid Similar to how invalid floating point arithmetic results in NaNs and not an exception. So, keep this in mind. Uh gotcha number four, non-array inputs. Again, this is added by design. It's not a bug. Um in NumPy, if we try to do a sum and we pass a Python list 1 2 3 here, we're going to get an expected result. So, that's going to be be six, I guess. Um in JAX, on the other hand, if we try to do this, we're going to get an exception. So, sum requires NumPy array or scalar arguments, got list at position zero. So, why is that? Uh let's imagine we try and make this uh more permissive. So, if we pass a Python list, we're going to convert it here to a JAX array, and now JAX sum will work on the JAX array. So, this thing will work. So, now let's create a list. X is a simple list. Now, let's try to using uh make JAX part to understand what will happen if we called jit on this function. So, let's try and run this. And you'll see that uh we'll basically have this whole thing unrolled, which is super inefficient. So, um the thing that happens is that we'll be passing element by element from this list, and JAX will not be able to optimize this and create some smart to use some smart primitives for looping, and instead we're going to get this inefficient optimization. So, yeah, uh keep this in mind. Now, this is a a really good place where upon which I can base my conclusion that JAX is really nice for researchers, but if you're a beginner, uh I don't think this will be I don't think that optimization details should be more important than the ease of use, such as the fact that you will not have an exception if you pass in a Python list instead of like an umpire or a JAX array. So, definitely things like this make me uh if I were to recommend to a beginner, uh I would definitely still recommend PyTorch, to be honest, uh over over JAX. Uh but like if you're a researcher and you want to have a lot of flexibility and you want to have like a super optimized uh programs, I think JAX is a really good uh bet for those guys. Anyways, let's continue. Uh gotcha number five, random numbers. I mentioned this already. Um now, let's see it in a bit more detail. So, basically what we can see here is that uh NumPy has a stateful uh pseudo-random number generator. That means that if I execute these two functions in a row, we'll have So, this one will uh advance the state of the uh PRNG, and this one after execution will also advance it, and it's hidden from us. In order to understand this a bit better, let's let's kind of dissect the the the this generator. So, we set the seed uh to some value. Seed is I think set to zero here, and basically we can fetch There is this function get_state, we can fetch the actual state of the NumPy uh PRNG. And uh I'm going to print certain um metadata from that state. Uh then I'm going to execute call like a sample uh a single number, and then we're going to fetch the state again, and we're going to sample the number again here, and we're going to fetch the state again and print it. So, let's see what happens uh here. So, first thing you can notice here is that the numbers here are different because the state is is internally uh advanced. So, what you can notice here is that first of all, uh NumPy is using this uh thing called Mersenne Twister. Hopefully, I'm pronouncing it correctly. Uh PRNG, and it's known to have a number of problems. Uh there is a link in their documentation for why that is. It's not that important for us. You know, what you probably useful to know is that it's basically has the state is consists of uh 624 uh like unsigned integers, and every time you call uh you sample from that generator, uh the we're basically uh consuming the the entropy of the generator. Bottom line, what I want you to take out from this is that uh NumPy's PRNG uh has some problems, and it's stateful, which is problematic because remember we need to have pure functions. We are in the functional programming uh like a world. On the other hand, this is how JAX operates. So, basically, you have the PRNG from JAX. You seed it with a certain value, and it gives back the key, and key is just a synonym for state. So, basically, we're now manipulating the state of the PRNG externally as opposed to internally. So, uh key is simply a like a tuple of two unsigned integers, uh 32-bit integers. So, that's the the state. Now, what's the trick here? So, now if we were to sample uh using that same state, we're going to get obviously uh the same results. So, that's different compared to the NumPy behavior, uh and the reason is again because we are not modifying the state, we're not advancing the state, uh and that's why we always get the same result. So, again, important to to notice here, state is preserved, state has not changed, and secondly, the results are the same. So, what do we do? How do we uh like we obviously this is not a random number, this is a constant function. So, how do we get a randomness in in in JAX? So, here is a trick. Every time you want to create a new random value, you basically just call the split function, and it's going to return uh key and subkey. Uh subkey is used to generate a novel number, and key can subsequently be used again in the split function to get a novel key and novel subkeys. So, every time you want to generate new number, you have to call this splitting. So, this may seem uh like very rough and problematic, uh but believe me, it helps uh solve various uh issues that are caused by randomness in in in libraries such as NumPy, etc., when you try to use them out outside of the context for which they were designed for, and that's like basically, I guess, CPU uh single-threaded programs, etc. Okay, so let's run this cell now and see the results. So, we have the old key. Uh old key got converted into the new key, so that's the new state, uh and we have this subkey, uh which we used to generate the uh random number. So, a couple of notes here. First things first is that um basically, you can split into more subkeys than just the two of these. And secondly, uh there is no semantic difference between uh the like the key and the subkey. These are basically recommendations for how to organize these uh states, and basically you you use key to generate novel to split and generate novel key and novel subkey, and you use the subkey to generate the current random number that you currently need to consume. And that's pretty much it. So, uh after having uh explained all these uh complex details about how to handle PRNGs compared to NumPy, is there any good reason for it for it? And the answer is yes. So, why this design? And the answer is um Can the code with the current design with NumPy's design be uh reproducible, parallelizable, and vectorizable? In the case of NumPy, number one is uh obeyed uh in the case where you have a single-threaded uh program on the CPU, uh basically there there is no problem. So, let's let's see this this concrete example. Let's see Let's assume we have a function called bar uh that generates a random number, and a function called baz uh that also generates uh like a random number. Finally, we have a function foo, and I don't know if I'm pronouncing these right. So, this function returns uh bar plus two times uh baz. Uh and uh don't know if you can see the problem with this with this code once we start and try to call uh this foo foo function on on line 14. Now, the the problem is um NumPy assumes uh a single-threaded environment, and uh basically Python um guarantees that this will be executed, to the best of my knowledge, from left to right, which means every time we run this, uh we are going to get the same result back, which means uh number one is obeyed, so that means we have reproducible programs. So, now what happens if we jit this function, or if we um and the jit decides to uh basically parallelize the this this program and calls bar on one accelerator on one core and baz on another one. So, what can happen is that the order of execution can change, and if that happens, we may get different results. So, if the first function to be called returns uh 0.3 and the second one returns 0.4, uh depending the of the order, we'll either have 0.3 plus two times 0.4, or we'll have or we'll have uh basically 0.4 plus two times 0.3, and that basically leads to different outcomes, which means this will not be uh reproducible result in the case of parallel uh like uh computing on parallel like uh cores, machines, whatever. Similarly, for the vectorization uh part, basically, uh NumPy guarantees that if you generate uh like numbers, if you do this, like if you iterate uh in a for loop, and you generate uh sequentially random numbers, uh that will give you the same results as if you were to just set size to three. On the other hand, JAX does it differently. Uh if you generate individually, uh that will give you different results compared to if you generate from from uh all the numbers at once using a specific key. So, and here again, you can see the an example where we we generate three subkeys uh using the split function and not only two keys, as I previously uh mentioned that. Um So, let's run this function, and you can see that basically uh NumPy uh has the same results, whereas uh like uh JAX does not have the same results. So, if you're using the common SIMD pattern, so that's the I guess single instruction multiple data, which is common in machine learning, uh whereas you want to apply the same functionality across different batches, uh you sometimes want to have the same random randomness being applied across the the all of the batches, and uh NumPy does not allow to do that. That's how I understood this this this part about the violating the vectorization problem. Let me know if I if I got this wrong. Okay, that's that's pretty much it. That was the random numbers. We have just two more gotchas to go and that's it. So gotcha number four is the control flow. It's fairly we we've seen something similar to this. So basically control flow plus grad nothing is there is no problem. Basically this transform function so the grad transform function can deal with these types of conditioning on the value of a function. I went ahead and ran this and you can see the results. Basically yeah, just analyze this function you'll you'll you'll see that at three there is this this jump where we have piecewise defined function here. But the whole point is if we were to take a grad of this function and evaluated at two and at four so that's just before this discontinuity and just after it we'll get valid valid gradients. Let me run this again. Just to update the state of the cell and if we were to jit this function and basically it will fail because we're conditioning on X. So we have to the solution is to to to make it a static argument and then it's going to work. So if I run this it's going to work and we already saw how to handle this case of conditioning on the value. So let's see some more interesting cases. So here we are conditioning again on the value but this time the value decides the length the the number of loops will have in this in this function. And the way to go around this is again make this a static static argument which will make jit trace this this function using the for for X using the abstract tracer so the shape and data type whereas for N is going to use a concrete value. So let me run this cell and see what we get. Okay, so as you can see we have some huge huge JAX pr here and the reason being is because we have a as you can see here 15 was passed for N and this is the best way that jit can can deal with these types of of primitive native Python for loops. Again importantly uh you should not change static values too too often otherwise we'll be triggering recompilation all of the time and then the overhead will maybe be detrimental to the speed of your application. Let's see how this can be avoided and make a bit more optimal. So a better way to do this is to use the low level API again so the lax API and there's this as we saw for I loop function and we can rewrite the above problem and I went ahead and just rewrote it here and you can see that um so this is the same function as the above one just using the lax API and if we now do the make if I call the make JAX pr and let me let me kind of run this we'll see that this code is way more succinct it's more concise compared to the above one and thus more efficient I guess. I haven't profiled this example but yeah I can assume it's more efficient. You can go ahead and analyze this these two cells at your own pace but the whole point was to be aware just be aware that using these lower level API functions you can sometimes make get more out of jit. Finally the only reason I have this one is to understand that you can condition sometimes on the dimensionality of your data and that's allowed because imagine this means that for the whole class of X where you have where X is two dimensional this can like a function can be cached and will work and we can fetch it and use it. So let's run this one and let's see the results. So I passed like a it is an input array which is not two dimensional and because of that we took this branch which means whatever we get as an input we just return as the output and that's why we have this super simple like a JAX pr for this particular function being traced with this input. Hopefully all of these helped you crystallize and understand jit because I guess I I think that's arguably the the hardest part to understand about JAX how jit works in the background. That's it. Final final final gotcha is like a how to handle NaNs in JAX. The usual non-error behavior is to simply return a NaN in the case of like operations such as this one division by zero. So if I were to comment out this thing here and let me run this cell so we we won't have any error we'll just have a NaN. So if you want to actually debug and understand where the NaNs came from you want to like the program to throw some exceptions you can do for example this you can find more in the docs but basically there is a there is a way to do it and now it will be throwing exceptions I guess. Let me try again. Yeah, now it's throwing an exception. So yeah, that's keep that in mind. Final cell JAX enforces single precision why because actually nowadays it's fairly common to train your models especially the big models like big transformers in FP16 or mixed precision or even FP8 so means eight bits only and because NumPy is aggressive in promoting like certain variables into double so that means 64 bits JAX by design that they are enforcing 32 bits. So that may lead to some again some some peculiar behavior such as this one. You say I want a like a a vector with thousand random numbers which are 64 like bit long and if we were to print the data type of the actual array we'll have float 32 which is not intuitive right? So be aware of this and there is a way around it you can set certain flags if you don't want to have this as a default behavior. As a quick summary we saw the basics of JAX such as the the the fact that it's using is is based on functional programming paradigm. We saw like various transform functions such as jit grad vmap. Then we saw we went deeper into into JAX. We saw the layered like onion like layered API. We saw many details of how jit works and we saw many gotchas that JAX has which may catch you by surprise so you should be aware and cognizant of these. Finally some conclusion of mine to who should be using JAX and who should be using other frameworks. As I said I think JAX is very good if you're a researcher you want to be very flexible and you want to have a powerful tool and you also want to once you train big models such as in whatever industry lab you're working you're going to be training big models so having like very optimized code is super important. But if you're a beginner on the other hand I think it may be a bit too harsh for you currently with the current state of JAX and I don't think that's going to change. So basically all of these optimization details such as the fact you cannot use Python list that will just throw exception because it can hurt the performance. Because of all of that and because of the functional paradigm it's I think it's hard it's easier to just start with PyTorch especially when we get to neural networks it's a more natural approach to approach neural networks with object-oriented from the object-oriented perspective at least in my opinion maybe I lack experience with JAX so I but that seems to be the case. Anyways in the next video we're going to cover some concepts such as pytrees handling states which are all basic components that we'll need to later build neural networks from scratch in pure JAX and also in Haiku or or Flax. Hopefully you found this video useful it took a lot of time to prepare the notebook to prepare everything that I wanted to show you here. So if you did like it consider sharing it out. Also consider subscribing to this channel if you haven't already. Also do join the discord community there is a lot of smart people there and we have a community for like