Machine Learning with JAX - From Zero to Hero | Tutorial #1
❤️ Become The AI Epiphany Patreon ❤️
https://www.patreon.com/theaiepiphany

👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦
https://discord.gg/peBrCpheKE

With this video I'm kicking off a series of tutorials on JAX! 

JAX is a powerful and increasingly more popular ML library built by the Google Research team. The 2 most popular deep learning frameworks built on top of JAX are Haiku (DeepMInd) and Flax (Google Research).

In this video I cover the basics as well as the nitty-gritty details of jit, grad, vmap, and various other idiosyncrasies of JAX.

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ JAX GitHub: https://github.com/google/jax
✅ JAX docs: https://jax.readthedocs.io/

✅ My notebook: https://github.com/gordicaleksa/get-started-with-JAX
✅ Useful video on autodiff: https://www.youtube.com/watch?v=wG_nF1awSSY
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable: 
00:00:00 What is JAX? JAX ecosystem
00:03:35 JAX basics
00:10:05 JAX is accelerator agnostic
00:15:00 jit explained
00:17:45 grad explained
00:27:25 The power of JAX autodiff (Hessians and beyond)
00:31:00 vmap explained
00:36:50 JAX API (NumPy, lax, XLA)
00:39:40 The nitty-gritty details of jit
00:46:55 Static arguments
00:50:05 Gotcha 1: Pure functions
00:56:00 Gotcha 2: In-Place Updates
00:57:35 Gotcha 3: Out-of-Bounds Indexing
00:59:55 Gotcha 4: Non-Array Inputs
01:01:50 Gotcha 5: Random Numbers
01:09:40 Gotcha 6: Control Flow
01:13:45 Gotcha 7: NaNs and float32
02:15:25 Quick summary
02:16:00 Conclusion: who should be using JAX?
02:17:10 Outro

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ 

If these videos, GitHub projects, and blogs help you,
consider helping me out by supporting me on Patreon! 

The AI Epiphany - https://www.patreon.com/theaiepiphany
One-time donation - https://www.paypal.com/paypalme/theaiepiphany 

Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković
Bartłomiej Danek
Zvonimir Sabljic

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

💼 LinkedIn - https://www.linkedin.com/in/aleksagordic/ 
🐦 Twitter - https://twitter.com/gordic_aleksa 
👨‍👩‍👧‍👦 Discord - https://discord.gg/peBrCpheKE

📺 YouTube - https://www.youtube.com/c/TheAIEpiphany/
📚 Medium - https://gordicaleksa.medium.com/ 
💻 GitHub - https://github.com/gordicaleksa 
📢 AI Newsletter - https://aiepiphany.substack.com/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#jax #machinelearning #framework
 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 box here 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 gonna see a lot of XLA later. So I thought just kind of explaining the terminology So they're 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 flex 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 graphML or RL they've been developing like a purposeful a Dedicated library built on top of Jax. So for example for graphML They've built up a thing called GRF or Jax for graphs. That would be like a acronym pretty much having mentioned flex 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 flex and that's the reason why I'll be covering both Haiku and flex 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 neural networks from scratch in both pure Jax as well as flex 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 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 higher levels of abstraction And as I said the last three videos there 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 That there is a reason behind that first if I was writing if I was coding everything from scratch That would take this video would be two 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 somewhat 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 I enabled the GPU accelerator in the background and The this Jupiter file which I'm gonna share with you after this video I'm gonna open source into 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 go 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 gonna 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 off 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 are importing from Jax There is this module called numPy and the convention is just to import it as J and P They also have a scipy API so Jax scipy, but we're not going to use that one in this view So these functions that are called the transform functions are a vital component of Jax So those are grad JIT V map and P map and you'll be seeing these a lot throughout this this notebook and in general using Jax so aside from that Jax's structure so the API structures like an onion as Usually software like libraries are so we have the high level API, which is a numPy Then we have this thing called called Lex and finally we have I guess XLA, which is I think in simple spots, right? Okay, so because 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 thousand points on the x-axis equally spaced from 0 to 10 And I just plot the function to sine cosine and we just plotted 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 matplotlibs 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 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 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 functional programs, so what that means is the following 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 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 are going to get an exception here as you can see Basically, Jax is complaining that we it says here object does not support item assignment. So you cannot Basically basically modify the array in place. So this is a 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 this in place and avoid allocating another memory object. So yeah, that's taken care of So that's the the second thing I wanted you to understand So the functional programming paradigm is that you can use the same array in place 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 moment But the key point here is you have to create this basically key which is a fancy name for like a state So whereas NumPy's pseudo random number generators are stateful, Jax is not stateful because it's As I said, that's a consequence of functional programming paradigm again So key is just a synonym 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 a trick And that's why Jax code is accelerator agnostic Basically, Jax automatically puts this array onto the accelerator So in this case, I'm using here my accelerator is set up 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 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 naturally leads me to fact number four And that's that Jax's accelerator agnostic same code runs everywhere You do not have to write particular syntax for a particular accelerators That's very cool Especially in the future where we see more and more companies being like building custom AI Accelerators and chips, not just TPUs and GPUs But like companies such as Graphcore, Cerebras, etc Let's see it in practice So here we define a Jax array Again, we see the syntax where we have the key We have to pass explicitly the key or the state of the pseudo random number generator And as you can notice, there is some difference in the API design here We have D type here We're using for NumPy, we're using S type There are some minor differences But like in general, syntaxes are very, very similar between Jax's NumPy and NumPy API I mentioned that we don't have to use two device anymore So that's in PyTorch So this array here is directly pushed to the GPU unit Whereas this one is on the CPU If we do some time profiling here, we can see that if let me run this thing Let me run this thing, basically this first line, line number 10 We do like a basically matrix multiply directly on the GPU And that's very fast We can see that we're going to see results in a second The second line here, because we have a NumPy dot And because we have NumPy arrays which are on CPU It's going to be way slower And NumPy only works, the reason why Jax exists is because NumPy only works with CPUs And we do need to leverage existing accelerators The third line here does something like a modified version of the first two lines Basically, we're using Jax's dot product And we are using NumPy arrays which will cause this line to push the devices to the GPU And then do the dot product And because of that overhead, we're going to have a slower result Hopefully for this third line 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 97 milliseconds for this line here Because we have the overhead of sending the data from the CPU, from the host to the GPU unit And the final number 26 is the same as this one And that's this one here So basically what we do, we can explicitly, using this device put function We can push the DRA to the GPU And then now we have equivalent lines between this one and line number 10 And that's it, that's why we have the same results here A couple of notes here First, I'm using GPU as a synonym for AI accelerator In reality, as I said, Jax's accelerator agnostic Which means basically, depending on what I set here in the runtime setting here We are going to run this either on GPU or TPU or whatever they have supported A second note is this block until ready You can notice that every time I'm using Jax's dot function, I'm using this block until ready The reason is 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 If I were to run this and maybe like let me assign that to a variable And now I do some coding here, blah, blah, blah 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 variable, this temporary variable is still not filled in Let's put it that way That's a fairly neat feature that makes things 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 the time that's necessary to dispatch Just delegate the task to the accelerator Just be aware of that Nice, those were some basics Now we're going to cover the transform functions But before that, let me do a quick overview So Jax's AI accelerator agnostic Jax handles random numbers a bit differently We're going to see why that is important a bit later Basically reproducibility in the environments where you have multiple accelerators Where you're doing parallel programming is like possible with Jax Whereas NumPy has some problems there because it was designed mainly for CPU programs Then as I said, functional programming paradigm is something you need to get comfortable with And finally, basically the syntax is very similar to NumPy That's a high level overview of what Jax is And now we're going to get deeper into these transform functions So as I said, JIT compiles your functions using this XLA And basically doing that, it catches some functions and makes it very, very fast to execute our programs 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 It's a cellU function, it's just an activation function The details here are not that important That's why I have this visualize function here I'm just going to run it We're going to see how it looks like You can see here up until zero and that's defined here When X is greater than zero, we just have X Which means we have a linear function here And when the number X is smaller than zero, we just have X And smaller than zero, we have some combination of exponents here with some coefficients, etc But that's not important The important part here is we can transform cellU So we pass the function inside of this JIT 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 for the sake of argument, we just have a super fast function right now here And let's benchmark it So let me allocate a vector of million data points We have million random numbers here And we pass them through the cellU function And we basically do a time profiling of both the normal version as well as the JIT version Let me run it again and let's see the numbers So the thing here is this function is not overly complex So the results we get with JIT So we will not get a huge performance boost here But imagine once you start training neural networks Where JIT starts shining because then the optimizer, the compiler has a bigger freedom To do various, various optimization tricks such as fusing the operators Such as avoiding allocating certain temporary structures, etc And things get really, really fast like orders of magnitude faster Here we can see that basically we have 1.96 milliseconds using the normal version And then we have only 121 microseconds using the JIT version So even in this simple example, we get performance benefits So for now, I'm going to leave it here I'm going to stay like on a high level here We don't have to understand currently how JIT works But like you just need to know when you see JIT, JIT makes functions run fast Let me try to be this guy Awesome Okay, let's go ahead here Second transform function, very important, is grad Basically it does the magic of automatic differentiation for you So the same thing as dot backward in PyTorch if you're coming from the PyTorch world And like a short note here, differentiation can be manual So basically you have a function, you manually calculate on the paper How the derivatives look like And then you can code that knowledge into your function And that's a manual differentiation Symbolic one is very similar It's just an automatic way to automate the manual differentiation Whereby the program is using those rules such as product rule, etc To build up derivatives of a function And numeric functions are things like finite derivatives Where you use basically numeric methods to compute the derivatives And finally, the automatic one is the one we all love Which is used in every single deep learning framework we know of Okay, let's see how grad works So let's define a simple 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 I just input, I just create like array of three values here, 0, 1, 2 And I rename the function to loss because this could be used as a loss function Just giving some semantics And finally, we can do grad by just wrapping loss into grad Again, we are passing function into a function Which is something you see often in the functional program in paradigm That gives us the grad loss So the gradient of the loss function which we can evaluate So this is contrast this to PyTorch or TensorFlow Where you just do backward Here you actually get the function back And you can evaluate it at particular points And by default, I said here 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 gonna be x1 times 2 plus x2 times 2 plus x3 times 2 And if I were to run this, we're going to get 024 which makes sense because we have 012 Let me print this out So print x, so we have 012 and we get 024 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 2x1 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 are some small errors here which are a natural consequence of the fact that we're storing numbers using a finite number of bits So how this works is very simple We take, so remember x is this array here, 012 Length of x thus will return us 3 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 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 2 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 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 as a value 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 So at least in the first iteration we're going to evaluate it at 0 And then we're going to add this small epsilon here So epsilon and then we're going to have 1, 2 And we're going to evaluate the function here We are going to subtract from that x minus epsilon So sorry, so that's 0 minus epsilon This thing will be constant for the first iteration And then we divide this by 2 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-brooks as expected Now let's see some fun examples We're going to define a simple second order polynomial function here x squared plus x plus 4 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 It gives them this gut feeling and it's easier to understand what's going on We're going to do higher order derivatives here Just to show you how a powerful grad function actually is So you can do grad of f and then you can do grad 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 plus 1 The second derivative so derivative of this one will just leave us with 2 And finally derivative of a constant is 0 That's why we have the third derivative is equal 0 So if we print all the values here we expect to see because x is 1 We expect to see 3, 2 and 0, right? So here it is 3, 2 and 0 The first value is just the value of the function at the input x which is 1 Which is as you can see here approximately 6 Exactly 6, okay? 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 gonna see how powerful this whole autodiff package of JAX is a bit later And now I'm gonna 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 gonna skip the evaluation because the evaluation will fail now And now I'm going to do the following If we just do the grads of the function By default do the derivative with respect to the first parameter That means we're going to get this thing here exactly So we should expect exactly the same numbers Except for the fact that f will be different when we evaluate it at x So this obviously will not work until I enter some numbers So this is going to be like let me say it's gonna be 1 And then I'm gonna modify this here So 1 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 7 because once you evaluate it at y equals 1 This will be 6 plus 1 that's 7 So how can we do derivative with respect to y? The only thing we need to do here is arc nums equals 1 instead I think this is the syntax Let me try it out Let me just ignore this for the time being I'm just gonna delete this And let me print this out So we get the number 2 here Which is correct result If we do a differentiation of this polynomial with respect to y We expect to get 2 times y Since y is equal 1 we get 2 here So the thing that confuses me here is I'm not sure why we are getting device array here Whereas previously we had pure floats on the CPU I'm not sure about that Let me check if I can do it like this Yep, for some reason once I add a comma here It's returning me device array Whatever Obviously this is now df dy But yeah, you get the point Now let me show you how powerful the autodef engine is So aside from grad which can do derivatives with respect for the scalar output functions Jacobians, so these jack forward and jack rev functions can find Jacobians So Jacobians basically can evaluate derivatives even for vector-valued functions That's the only difference That's pretty much the only difference So here let me take another function Basically we have a simple paraboloid If we google it here we can see how it looks like It looks like a cup 2D 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 figure out the second order derivatives We can get the Hessian number So we can get the second order derivative with respect to x is 2 As you can see here for y it's also going to be 2 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 a collection It's a matrix of as you can see here various derivatives of the multivariate function You can simply define it using the jack rev and jack forward function We additionally did 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 jack forward works 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 2, 2 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 I just wanted to show you that this is possible And if you want to dig into more details there documentation is very nice And yeah let's continue with the final example for grad I just took this edge case function so the absolute value of x And let's see how how how Jax handles it Basically this is how the function looks like I just visualized it here And I printed the values at minus 1 and 1 which would be 1 and 1 And you can see the the results here And finally I printed the I found the gradient of this function And you can see it's not differentiable at this point at 0 It's the gradient is undefined But like if we were to evaluate the function at minus 1 we get So as you can see here the derivative will be minus 1 for all the numbers here 1 for all numbers here The interesting point is actually 0 Because as you can see this one returns the obvious result minus 1 So this one here is kind of interesting If we were to delete this part and just kind of add a small number to 0 We'll get 1 as the output If we were to do the same thing just on the negative side we're going to get minus 1 So basically what they've done is they defined for 0 is going to be the value is going to be defined is going to be 1 That's just a convention and that's how Jax deals with it We didn't get any exception so yeah I guess I guess that's good at least in some settings Finally let's jump to Vmap Basically when you see Vmap you need to think about it the main value prop of Vmap is the following Writer functions as if you were dealing with a single data point And this is going to become more clear as we go through these examples Let's say we have a matrix W which is basically weights of you can treat it as weights of a linear neural network layer We input some state we set the shape to 150, 100 And we now create this batch X 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 times 10 pixel image it's 100 when you once you flattened out So now if we apply the matrix we do the dot product between W and X This is basically a simulating what the linear layer is doing in the background when it's processing the input data Now the trick here is this only transforms the single a single image it cannot handle the batch Otherwise it will crash because we're trying to multiply 150, 100 which is W And we're trying to multiply that with 10, 100 that's going to fail So this thing does not work for the batch so okay let me run it So how would we go about making this work for a batch of images? So that's because that's what we care about and because that's way more efficient as you may know Especially on GPUs and accelerators Basically the most naive approach would be to iterate through the images and then call this function that knows how to handle single data points And then we just stack the results and basically we can see that this will work And the thing is it's very very slow because you want to avoid doing for loops you want to vectorize your functions And that's why this runs 5.54 milliseconds Now let's see a bit more a better approach of doing this So this would be a better approach but as you can see we had to completely rewrite the function 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 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 what we expected right? So everything is fine now the problem with this and it's additional jitted so it's going to be really fast Now the problem with this is as you can see we have to write depending on whether we handle singular cases 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 Jacks offers us using something called V map You just take the function that handles a single data point You wrap it into this you transform it using this V map And you can now pass the batch of your data without any other modifications So that's that's very cool. If I run this Let's see how how fast this thing is. So it's around 100 microseconds I think the last time I ran this this was actually faster than this one But I guess there are some some variants 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 Functions and it's as as efficient as by like doing this laborious work So basically what what V map does in the background is it takes the the for loops and packs them into this Alex API mid-level layer of the of the API and yeah, so now let me let me go ahead and 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 I guess in in line with the functional 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 as the batch of data so We'll need to add this thing here again. So if I were to run this right now, it will crash and we'll see why so it basically says It basically says that W We're trying to to tell to V map that W has a batch dimension where whereas it does not have it's just a it's just a matrix 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 And yeah, it did work. So basically that's it Now you saw in a bit more detail how V map 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 V map and now let's dig Even deeper and understand the intricacies of how JIT works because that will help you enormously Debugging these Jax programs. So let's let's continue here. So We basically have as I said Jax has this onion like API layer API structure and that's that's I guess pretty much Always the case but still Anyways, we have numpy as the as the highest level then we have lax as the mid level and finally the XLA 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 in the numpy level of the API This thing can be tolerated So one plus one like one as the integer plus one as the integer Like one as the integer plus one as the float will work. But once we get to the mid level Here we need to be explicit about the types. Otherwise we'll have we'll get an error So if I run this we'll see that this is printed So one plus one is okay, but here add requires arguments to have the same D types got int 32 and float 32 The reason they've done this done it this way is to be more error robust Finally, we have the fact that the lax layer is obviously more powerful Although as a trade-off, it's it's less user friendly, which is kind of obvious as an example here. We have X and Y We have basically we want to do a convolution a 1d convolution between the signals X and Y and in numpy API this would be done like this We just call the convol function and we get the result On the other hand, once you get into the lax land, you have to use this conv general dilated Which is way more powerful as you can see 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 we can see the results here as well as here will be the same Like this this assert should make sure that that we get the same results And as you can see it works perfectly now Final remark here is you can see here that this result returns batched results So that's why we have to index 0 0 to get the actual result As I said, lax API is just a thin wrapper around thin I mean, it's a wrapper around XLA so you can actually find the the XLA function that's going to be eventually called here And that's this conv the general padding It's in C++ you can see the arguments etc So if you ever need to dig a bit deeper and optimize something really really really hard Then you TensorFlow documentation here got you covered So let's get back here and continue That was the short mention of the API Now let's finally understand how it works So again, I won't get into more details here The whole point of this is to show that JIT functions are faster compared to non-JIT versions of the function Here we just normalize the matrix X like 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, the difference here is again not that big because this is a super simple function The more complex the function is the bigger the differences here will be between the JIT version and between the normal version So here we have the JIT version and again we're using a block until ready because of the asynchronous dispatch and we see the results Okay, so let's see In order to understand JIT it may be useful to understand when it fails Which functions cannot, so which class of functions we cannot JIT will tell you, will help you understand how JIT actually works behind the curtains So if I were to run this thing it's gonna crash So we are creating a simple vector of 10 random elements and we pass that vector into this get negatives and we call it So now because we don't have any JIT we're gonna get a result back and everything works as expected But if I were to call the JIT version of the function by wrapping this into this JIT transform If we run this we're going to get like error and 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 We are going to get an output that's going to vary in its shape and that's something that's not 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 like a dot product between X and Y and returns result and in the meanwhile it also prints some intermediate variables here So I prepare the variable X and Y, just random vectors and matrix here and we call the function here And then we again call, we have a second call here and let's see what happens after the second call So I'm gonna run this, there's a couple things happening here, first thing first is the first time you run a JIT function So that's line 14, what JIT does in the background is something called tracing So it basically takes, instead of inputting the actual values of X and Y, it creates this abstract like 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 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 sampled from Gaussian, we have the shape here 3, 4 and we have 4 for Y So basically you can see these are actually passed into the function the first time JIT is called And using that JIT can understand how the shapes are morphing 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 the shape is 3 We finally get the results here, the second time you call the function something funny happens and that has to do with the functional programming paradigm Basically because print functions are a side effect because we are returning the result from this function by other means than through the actual output here So we are printing and that's a side effect and because of that 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 compile function, we won't see any printing And the whole point is this caching mechanism I just explained, so that's why JIT functions are so fast So I mentioned here anytime we get the same shapes and types we just call the compile function So that means if I were to now call whatever, so basically for the whole class of inputs that have this shape and this data type And no matter the value we are always going to call the compile function But if I were to call this function again but this time with a different shape, so let's do something like this So x3 maybe 3, 5 and this will be 5 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 we'll again trigger the compilation because the shape changed here And so JIT is smart enough to retrace it and as you can see here we have the tracing happen here again Hopefully you get a better picture of how JIT now works but we are going to continue and understand this in even more depth So now 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 functions when you want to use them with JAKS transform functions such as JIT So no side effects If I were to print this and we are going to have something called JAKSper here, so that's JAKS expression And that's basically the, let's call it like a flow model that JIT creates in the background when it does its tracing procedure So if I were to run this, let's see what we have here So you can see what happens, so it creates this abstract grammar and you basically have C, so add 1 to A So that's going to be, 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 B plus 1, so that's Y plus 1, this is going to be B And then it calls 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 in more detail every single part of this syntax of this grammar But like, now you understand a bit better how JIT, when it does the tracing, creates this type of a flow in the background And this can be compiled using XLA, okay? Let's see another example of a failure This time, what we do is we pass this argument neck and we condition upon it And depending on the value, whether it's true or false, we're going to return minus X or plus X And remember, the first time we call a JIT function, it's going to input abstract shape and like a data type It won't have information about the value, and that's why this thing is going to fail So it says here, blah blah blah, abstract tracer, value encountered, where concrete value is expected The problem arose in the bool function, so that's basically what happens And in order to avoid this, what we can do is use these static arguments By making an argument static, so we say here, hey, first argument, so that's neck, is going to be static What we do 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 value So we are kind of lowering the level of abstraction here while doing the tracing Which is going to kind of constrain the class of inputs where this compile function can be called from the cache But in return, yeah, in return, we get it to work, I guess Okay, so let's call this thing, so this should now work And the thing I want you to understand here again is that the first time we call it with true here, we do the tracing And then the second time we call it with true, because nothing changed, 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 And we'll have two cached functions, so that's like, which run really fast Let's continue analyzing the failures, the third failure is here If I were to run this, let's see what we get, 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 x shape, let's print this whole product thing And I think that's going to work, so 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 trace object, and we're trying to pass the trace object into reshape, which expects a concrete value That's why this thing is crashing again Let's see what the solution is, we can basically use NumPy prod instead of Jack's product So this may be a little bit confusing, basically as you saw, you have two tools to make these functions work legit One is to make certain arguments static, and sometimes you'll have to use NumPy functions instead of Jack's functions So don't ask me why, I'm only a couple of steps ahead of you, I recently started learning Jack's myself But let's try and run this, and this time it works, because we have the mprod will return a concrete value and not the traced object And that's it, if you wish, as I said, I'm going to share the colab, I'm going to push it to my github, so go ahead and play with these failure cases yourself To better understand how JIT works, but basically just keep in mind that JIT passes in these abstract objects that have the shape information and the data type information and no value And that will help you save yourself from a lot of headache, I guess Ah, okay, this is something I need to cover because there are some gotchas, some idiosyncrasies of Jack's 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 functions, so Jack's is designed to work only on pure functions, we saw some glimpse of this while we were using print functions a couple of 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 the inputs, you just retrieve the outputs So that's an informal definition 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 this return statement here So let me call this function, okay, on the first call, again, we have the tracing, so we call the executing function print statement, but the second call, due to the fact that we are passing the same shape and same data type, that's just a float 32 here We just use the cache function and we return five, finally, the third call, because we changed the shape and the type here, so we basically have a ray now, it's going to basically call and trigger 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 program 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 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 what will happen here, once we call the jitted version of the function the first time, it's going to cache the value of g, which is zero Then we are going to update it, and then we are going to call it with five, and it's going to call the cached version of the function, which, if you recall, has g equals to zero somewhere in the Jackspur, and that's why we are going to get the wrong results So let me execute this cell, so on the first call we get four, and that's fine, because we passed four, we had g equals zero, that's fine, now g is equal to ten, and we have five here, so we expect fifteen, but we get five because the g version of zero is cached in the function And finally, if I call with a different, we change, instead of float, we pass the, like, a Jacks array, this time we trigger the execution and we get the correct result again So in any case, this is a wrong idea, a bad idea to add these types of impurities Example number three, so Haiku Flux are basically built upon this idea, the whole idea is that it's fine to have a stateful functionality inside of your pure function So what I mean by that is the following, so we have x here, and what we do here, 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 pure, because we're only using x, we're only assigning x to this state And we're only outputting whatever came out from this function, so we're not accessing some global variable or whatnot So because of that, all of this will be fine, and if I execute this, we're 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 use, 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 fail, even though we semantically would expect the same results, and you can see that's the case, we got 45 here, we got 0 here, even though we should have gotten 45 as well Here you can see this lex primitive, so that's one of those mid-level API functions, it's called foriloop, 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 understand here is just that you cannot use iterators because they are stateful, and we are thus violating the purity constraint of Jax Okay, that was the gotcha number one, make sure to write pure functions if you want to use them with Jax transforms in general Gotcha number two, in-place updates, we saw this already, so you cannot modify the arrays in place, you have to use that at set syntax, so let's see a simple example here Once we create the Jax array, we have to use this at set syntax in order to get 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 if you're not using this input array, then it will not allocate 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 do so on this higher level perspective Don't worry about the expressiveness, we can still do everything we can with NumPy, so here if I create a simple matrix of ones, we can do whatever we want, we can add, we don't just have to assign values 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 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 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 ten elements, so from zero to nine, and we try to index into eleventh position, this will throw an exception because there is no eleventh 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 this is out of bounds for x is zero with size ten Okay, let's see what Jax's behavior here is, if we were to try to assign at position eleven, so that's the same example as here, if we try to add twenty-three, it's just going to ignore 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, if we try to return, to retrieve the element at eleventh position, which does not exist, it's going to clamp eleventh to the last 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 information from you, but as a consequence you get this And I guess if you're familiar with NAND behavior, this is kind of similar, because NANDs are also in a sense not acceptable, although we also have non-error behavior, whereas we just get NANDs back instead of the system throwing some exception So I mentioned it here, similarly to how invalid floating point arithmetic results in NANDs and not an exception, so keep this in mind Gotcha number four, non-array inputs, again this is added by design, it's not a bug, in NumPy if we try to do a sum and we pass a Python list, one, two, three here, we are going to get an expected result, so that's going to be six I guess In Jax on the other hand, if we try to do this, we're going to get an exception, so sum requires a NumPy array or scalar arguments got list at position zero So why is that? Let's imagine we try and make this 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 make Jaxper to understand what will happen if we called Jit on this function, so let's try and run this, and you'll see that we'll basically have this whole thing unrolled, which is super inefficient So 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, keep this in mind, now this is 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, 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 a NumPy or Jax array So definitely things like this make me, if I were to recommend to a beginner, I would definitely still recommend PyTorch to be honest over Jax, 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 programs, I think Jax is a really good bet for those guys Anyways, let's continue, gotcha number five, random numbers, I mentioned this already, now let's see it in a bit more detail, so basically what we can see here is that NumPy has a stateful pseudo random number generator That means that if I execute these two functions in a row, we'll have, so this one will advance the state of the 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 kind of dissect this generator So we set the seed to some value, seed is I think set to zero here, and basically we can fetch, there is this function getState, we can fetch the actual state of the NumPy PRNG, and I'm gonna print certain metadata from that state Then I'm going to execute, like a sample, a single number, and then we're gonna fetch the state again, and we're gonna sample the number again here, and we're gonna fetch the state again and print it, so let's see what happens here So first thing you can notice here is that the numbers here are different because the state is internally advanced So what you can notice here is that first of all NumPy is using this thing called Mersenne Twister, hopefully I'm pronouncing it correctly, PRNG, and it's known to have a number of problems, there is a link in their documentation for why that is, it's not that important for us What is probably useful to know is that it basically has the state consists of 624 unsigned integers, and every time you sample from that generator, we are basically consuming the entropy of the generator Bottom line, what I want you to take out from this is that NumPy's PRNG has some problems, and it's stateful, which is problematic because remember we need to have pure functions, we are in the functional programming 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 are now manipulating the state of the PRNG externally as opposed to internally So key is simply like a tuple of two unsigned integers, 32-bit integers, so that's the state Now what's the trick here? So now if we were to sample using that same state, we are going to get obviously the same results, so that's different compared to the NumPy behavior And the reason is again because we are not modifying the state, we are not advancing the state, and that's why we always get the same results So again, important to notice here, state is preserved, state has not changed, and secondly, the results are the same So what do we do? Obviously this is not a random number, this is a constant function, so how do we get a randomness in Jax? So here is a trick, every time you want to create a new random value, you basically just call this split function, and it's going to return key and subkey Subkey is used to generate a novel number, and key can subsequently be used again in the split function to get novel key and novel subkeys So every time you want to generate a new number, you have to call this splitting So this may seem like very rough and problematic, but believe me, it helps solve various issues that are caused by randomness in libraries such as NumPy, etc When you try to use them outside of the context for which they were designed for, and that's like basically I guess CPU, single threaded programs, etc Okay, so let's run this cell now and see the results. So we have the old key, old key got converted into the new key, so that's the new state And we have this subkey which we use to generate the random number So a couple of notes here, first things first is that basically you can split into more subkeys than just the two of these And secondly, there is no semantic difference between the key and the subkey, these are basically recommendations for how to organize these states And basically you use key to generate novels 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 after having explained all these complex details about how to handle PRNGs compared to NumPy, is there any good reason for it? And the answer is yes. So why this design? And the answer is can the code with the current design, with NumPy's design be reproducible, parallelizable, and vectorizable? So in the case of NumPy, number one is obeyed, in the case where you have a single threaded program on the CPU, basically there is no problem So let's see this concrete example. Let's assume we have a function called bar that generates a random number and a function called baz that also generates a random number Finally we have a function foe, and I don't know if I'm pronouncing these right, so this function returns bar plus two times baz And I don't know if you can see the problem with this code, once we start and try to call this foe function on line 14 Now the problem is NumPy assumes a single threaded environment, and basically Python guarantees that this will be executed, to the best of my knowledge, from left to right Which means every time we run this, we are going to get the same result back, which means number one is obeyed, so that means we have reproducible programs So now what happens if we get this function, or if we, and the JIT decides to basically parallelize 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 0.3 and the second one returns 0.4, depending on the order we'll either have 0.3 plus 2 times 0.4, or we'll have basically 0.4 plus 2 times 0.3 And that basically leads to different outcomes, which means this will not be a reproducible result in the case of parallel, like computing on parallel, like cores, machines, whatever Similarly for the vectorization part, basically NumPy guarantees that if you generate numbers, if you do this, if you iterate in a for loop and you generate sequentially random numbers That will give you the same results as if you were to just set size to 3, on the other hand, Jax does it differently, if you generate individually That will give you different results compared to if you generate from all the numbers at once using a specific key So, and here again you can see an example where we generate three subkeys using the split function and not only two keys, as I previously mentioned that So let's run this function and you can see that basically NumPy has the same results, whereas Jax does not have the same results So if you're using the common SIMD pattern, so that's the single instruction multiple data, which is common in machine learning Whereas you want to apply the same functionality across different batches, you sometimes want to have the same randomness being applied across all of the batches And NumPy does not allow to do that, that's how I understood this part about the violating the vectorization problem, let me know if I got this wrong Okay, 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've seen something similar to this, so basically control flow plus grad, nothing, 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 see that at three there is this jump where we have a piecewise defined function here The whole point is if we were to take a grad of this function and evaluate it at two and at four, so that's just before this discontinuity and just after it, we'll get 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 are conditioning on x So we have to, the solution is to make it a static argument and then it's gonna work, so if I run this, it's going to work And we already saw how to handle this case of conditioning on a 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 number of loops we'll have in this function And the way to go around this is again make this a static argument which will make jit trace this function using the forex, using the abstract tracer, so the shape and data type Whereas for n is gonna 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 checks per here and the reason being is because we have, as you can see here, 15 was passed for n And this is the best way that jit can deal with these types of primitive, of native Python for loops Again, importantly, you should not change static values 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 lex API, and there is this, as we saw, for iLoop function And we can rewrite the above problem and I went ahead and just rewrote it here and you can see that So this is the same function as the above one just using the lex API and if we now do the make, if I call the make checks per And 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 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 get more out of it 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 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 Jacksper for this particular function being traced with this input Hopefully all of these helped you crystallize and understand JIT because I guess, I think that's arguably the hardest part to understand about JAKS How JIT works in the background That's it, final, final, final gotcha is like how to handle NANDs in JAKS The usual non-error behavior is to simply return a NAND 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 won't have any error, we'll just have a NAND So if we want to actually debug and understand where the NANDs came from You wanna 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 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, let's keep that in mind Final cell, JAKS 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 it means 8 bits only And because NAMPA is aggressive in promoting like certain variables into double, so that means 64 bits JAKS made it by design that they are enforcing 32 bits So that may lead to some, again, some peculiar behavior such as this one You say I wanna like a vector with thousand random numbers which are 64 like bit long And if we were to print the data top of the actual array, we'll have flow 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 JAKS such as the fact that it's using, it's based on functional programming paradigm We saw like various transform functions such as JIT, GRAD, DMAP Then we saw, we went deeper into JAKS, we saw the layered, like, layer API We saw many details of how JIT works and we saw many gotchas that JAKS 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 JAKS and who should be using other frameworks As I said, I think JAKS is very good if you're a researcher, you wanna be very flexible and you wanna have a powerful tool And you also wanna, once you train big models such as in whatever industry I lab you're working, you're gonna 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 JAKS 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 I think 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 JAKS so yeah, 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 PureJAKS and also in Haiku or Flex Hopefully you found this video useful, it took a lot of time to prepare the notebook, to prepare everything 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 there is more than 1,000-100 people at this point of time So somebody will help you out if you have some doubt or questions In any case, until next time, bye bye!
[{"start": 0.0, "end": 7.96, "text": " What's cracking guys with this video? I'm kicking off a series of video tutorials about Jax. So what is Jax exactly?"}, {"start": 7.96, "end": 9.96, "text": " Let's see the official box here"}, {"start": 10.8, "end": 14.32, "text": " Basically Jax is a autograd and XLA which stands for"}, {"start": 15.120000000000001, "end": 17.12, "text": " accelerated linear algebra"}, {"start": 17.16, "end": 22.02, "text": " compiler which is basically a compiler developed by Google that co-evolved with the"}, {"start": 22.88, "end": 24.84, "text": " development of TPU units"}, {"start": 24.84, "end": 30.04, "text": " We're gonna see a lot of XLA later. So I thought just kind of explaining the terminology"}, {"start": 30.04, "end": 35.88, "text": " So they're basically brought together for high-performance numerical computing and machine learning research"}, {"start": 36.08, "end": 40.92, "text": " so in a nutshell we can treat Jax as a like a machine learning library and"}, {"start": 41.6, "end": 46.8, "text": " you need to be cautious not to compare Jax directly to pytorch or to tensorflow because"}, {"start": 47.16, "end": 52.879999999999995, "text": " Jax is kind of like a mid-level like library and people have been building on top of it"}, {"start": 52.88, "end": 56.400000000000006, "text": " So we have as for the high-level deep learning libraries"}, {"start": 56.400000000000006, "end": 60.800000000000004, "text": " we have the two most popular ones are Haiku coming from from DeepMind and"}, {"start": 61.040000000000006, "end": 64.16, "text": " We have a flex coming from the Google research team"}, {"start": 64.72, "end": 72.28, "text": " So basically aside from that DeepMind team has been developing for every single like domain particular domain such as graphML"}, {"start": 72.28, "end": 75.84, "text": " or RL they've been developing like a purposeful a"}, {"start": 76.80000000000001, "end": 80.08, "text": " Dedicated library built on top of Jax. So for example for graphML"}, {"start": 80.08, "end": 88.12, "text": " They've built up a thing called GRF or Jax for graphs. That would be like a acronym pretty much having mentioned flex"}, {"start": 88.12, "end": 91.16, "text": " I think it's worth noting that on hugging face hub"}, {"start": 91.16, "end": 96.0, "text": " You have more than 5,000 models at the at the time of this recording"}, {"start": 96.6, "end": 98.92, "text": " That you can use download and play with them"}, {"start": 98.92, "end": 103.75999999999999, "text": " And I think most of them were written in flex and that's the reason why I'll be covering both"}, {"start": 103.92, "end": 108.38, "text": " Haiku and flex in some of the future videos that out of the way"}, {"start": 108.38, "end": 111.88, "text": " Let me tell you what I'm going to cover in this series of videos"}, {"start": 111.88, "end": 118.22, "text": " So in the first two videos I'm going to explain the nitty-gritty details of Jax and then we're going to do"}, {"start": 118.66, "end": 124.74, "text": " We're going to code neural networks from scratch in both pure Jax as well as flex as well as Haiku"}, {"start": 124.89999999999999, "end": 129.34, "text": " also those of you who know me, you know that I'm a big fan of a top-down approach when teaching and"}, {"start": 129.9, "end": 130.74, "text": " Basically here"}, {"start": 130.74, "end": 136.46, "text": " I'm going to go a bit differently because the topic is kind of complex and the first two videos are going to go"}, {"start": 136.46, "end": 142.74, "text": " So the first one is going to go from mid-level of abstraction down to lower lower levels of abstraction"}, {"start": 143.18, "end": 146.82, "text": " And we're going to understand the nitty-gritty details of Jax in this video"}, {"start": 146.94, "end": 150.34, "text": " The second one will be mid-level to higher levels of abstraction"}, {"start": 150.34, "end": 156.66, "text": " And as I said the last three videos there will be hands-on and those are going to be on a high level of abstraction"}, {"start": 156.66, "end": 161.66, "text": " We won't be digging deeper into the primitives. We're just going to build from the components"}, {"start": 161.66, "end": 165.14000000000001, "text": " We learned in the first two videos. We're going to build neural networks from from scratch"}, {"start": 165.14, "end": 169.77999999999997, "text": " So aside from that as you can see here, I have a code skeleton already"}, {"start": 170.5, "end": 171.94, "text": " here and"}, {"start": 171.94, "end": 178.01999999999998, "text": " That there is a reason behind that first if I was writing if I was coding everything from scratch"}, {"start": 178.26, "end": 180.85999999999999, "text": " That would take this video would be two hours long"}, {"start": 181.42, "end": 188.2, "text": " Secondly, this is way more natural because you usually never ever build build stuff from your head from scratch"}, {"start": 188.2, "end": 194.51999999999998, "text": " You take you either take some snippets from the documentation from stack overflow or from somewhat somebody else's repo"}, {"start": 194.52, "end": 197.06, "text": " So it's very I want to stress the understanding here"}, {"start": 197.06, "end": 200.10000000000002, "text": " It's important that we understand the code so that we can modify it"}, {"start": 200.10000000000002, "end": 204.94, "text": " And I think it's less important for you to understand how to code this up from your head because nobody does that"}, {"start": 204.94, "end": 211.56, "text": " I mean sometimes I even like forget how to write down if name equals main in the main like file of a Python program"}, {"start": 211.56, "end": 216.66000000000003, "text": " So you just search for stuff if you don't understand how to write something down. Okay, let's dig down into the code"}, {"start": 217.10000000000002, "end": 219.66000000000003, "text": " first off I'm in Google Colab here I"}, {"start": 219.66, "end": 223.82, "text": " I enabled the GPU accelerator in the background and"}, {"start": 224.46, "end": 228.7, "text": " The this Jupiter file which I'm gonna share with you after this video"}, {"start": 228.7, "end": 233.78, "text": " I'm gonna open source into my github is structured into two like sections"}, {"start": 233.78, "end": 240.26, "text": " The first one is warming up with Jax basically we're on the mid level of abstraction pyramid. Let's go that way and then we're going deeper"}, {"start": 241.14, "end": 245.24, "text": " Understanding the nitty-gritty details behind many of Jax primitive"}, {"start": 245.24, "end": 252.20000000000002, "text": " Transform functions such as JIT grad etc. We're gonna see those in a second. Okay, let's let's let's start. So first things first"}, {"start": 252.56, "end": 253.96, "text": " I"}, {"start": 253.96, "end": 256.72, "text": " want you to understand that Jax is"}, {"start": 257.68, "end": 260.32, "text": " Jax's API is very similar to numPy's and"}, {"start": 260.96000000000004, "end": 267.14, "text": " Second thing I want you to take off from this section is that Jax code is accelerator agnostic and we're going to see what that means"}, {"start": 267.14, "end": 269.84000000000003, "text": " Exactly. Okay, let's let's dig into the actual code"}, {"start": 269.84000000000003, "end": 274.6, "text": " so as I said for the most part the syntax is the same as numPy's and"}, {"start": 274.6, "end": 277.20000000000005, "text": " You can see here. We are importing from Jax"}, {"start": 277.20000000000005, "end": 281.8, "text": " There is this module called numPy and the convention is just to import it as J and P"}, {"start": 282.1, "end": 287.56, "text": " They also have a scipy API so Jax scipy, but we're not going to use that one in this view"}, {"start": 288.32000000000005, "end": 293.46000000000004, "text": " So these functions that are called the transform functions are a vital component of Jax"}, {"start": 293.48, "end": 299.04, "text": " So those are grad JIT V map and P map and you'll be seeing these a lot"}, {"start": 299.48, "end": 303.06, "text": " throughout this this notebook and in general using Jax so"}, {"start": 303.06, "end": 308.06, "text": " aside from that Jax's structure so the API structures like an onion as"}, {"start": 308.06, "end": 313.3, "text": " Usually software like libraries are so we have the high level API, which is a numPy"}, {"start": 313.3, "end": 319.46, "text": " Then we have this thing called called Lex and finally we have I guess XLA, which is I think in simple spots, right?"}, {"start": 319.46, "end": 321.46, "text": " Okay, so because I'm a"}, {"start": 322.46, "end": 327.9, "text": " Etymology nerd. I just thought writing this down. Basically Lex is just an anagram for XLA"}, {"start": 327.9, "end": 333.38, "text": " Which is a compiler and I'm not completely sure how people came up with the name Jax"}, {"start": 333.38, "end": 336.29999999999995, "text": " So if anybody knows feel free to type it down in the comment section"}, {"start": 336.29999999999995, "end": 337.85999999999996, "text": " Let's see a couple of facts about Jax"}, {"start": 337.85999999999996, "end": 343.34, "text": " So the first thing I already mentioned the syntax is remarkably similar to numPy's if I run this we can see we have"}, {"start": 343.34, "end": 348.85999999999996, "text": " We have defined thousand points on the x-axis equally spaced from 0 to 10"}, {"start": 348.85999999999996, "end": 354.29999999999995, "text": " And I just plot the function to sine cosine and we just plotted we get the chart here"}, {"start": 354.3, "end": 361.06, "text": " You can see that Jax syntax is completely similar like it's actually same if you just switch MP for JMP"}, {"start": 361.06, "end": 364.06, "text": " You get linspace sine cosine plot"}, {"start": 364.06, "end": 367.58000000000004, "text": " You can like as you can see here the arrays can be directly"}, {"start": 368.06, "end": 375.58000000000004, "text": " Inserted into the matplotlibs plot function. So if I run this one we get the same results. Okay, so that's it"}, {"start": 375.58000000000004, "end": 377.58000000000004, "text": " So that's the first fact second fact"}, {"start": 378.22, "end": 382.66, "text": " So that's something you need to get comfortable with and that's the functional programming"}, {"start": 382.66, "end": 387.06, "text": " Paradigm, especially if you're only familiar with the object-oriented programming paradigm"}, {"start": 387.06, "end": 389.06, "text": " This is going to be a bit rough"}, {"start": 389.06, "end": 394.86, "text": " But like I I promise you there are some awesome benefits that functional programming brings with it. So"}, {"start": 395.70000000000005, "end": 401.94000000000005, "text": " Jax arrays are immutable. So that's one of the three like that's one of the common properties you see in functional programming"}, {"start": 402.46000000000004, "end": 406.34000000000003, "text": " in functional programs, so what that means is the following"}, {"start": 406.34, "end": 413.61999999999995, "text": " If we take a like a NumPy array here, so numbers from 0 to 9 we print them out and I then modify the"}, {"start": 414.09999999999997, "end": 419.5, "text": " The array at index 0 and I add some random number like 23 and if I run this"}, {"start": 419.97999999999996, "end": 420.58, "text": " Everything is fine"}, {"start": 420.58, "end": 427.14, "text": " We have the original array and we have the modified array and as you can see X is modified in place, which means"}, {"start": 427.78, "end": 432.09999999999997, "text": " The array is by definition mutable on the other hand if we take this"}, {"start": 432.1, "end": 439.06, "text": " The same program and run it here. So I create this time Jax array instead of NumPy array and I try to modify the"}, {"start": 439.54, "end": 444.06, "text": " Array in place. This is what's going to happen. So we are going to get an exception here as you can see"}, {"start": 444.90000000000003, "end": 451.78000000000003, "text": " Basically, Jax is complaining that we it says here object does not support item assignment. So you cannot"}, {"start": 452.58000000000004, "end": 457.78000000000003, "text": " Basically basically modify the array in place. So this is a solution. It's kind of rough"}, {"start": 457.78, "end": 464.73999999999995, "text": " But yeah, you have to do this at syntax with index and then set the value and this is going to work"}, {"start": 464.73999999999995, "end": 471.61999999999995, "text": " Finally, you can see that we are not like modifying X in place. We basically allocate additional object in the memory"}, {"start": 472.26, "end": 476.17999999999995, "text": " Called Y and now we have both the original array as well as the"}, {"start": 476.73999999999995, "end": 481.61999999999995, "text": " Novel array so you may think here. Well, this is kind of suboptimal, right?"}, {"start": 481.61999999999995, "end": 484.41999999999996, "text": " Well, the trick is and we're going to see that a bit later"}, {"start": 484.42, "end": 488.98, "text": " Like Jax uses something called JIT which uses something called XLA which"}, {"start": 489.54, "end": 494.1, "text": " Basically takes care of this stuff behind the curtains. So you don't have to worry about this"}, {"start": 494.1, "end": 497.06, "text": " So if it notices that you're not using the original array"}, {"start": 497.06, "end": 503.86, "text": " It's going to just do this in place and avoid allocating another memory object. So yeah, that's taken care of"}, {"start": 504.66, "end": 507.78000000000003, "text": " So that's the the second thing I wanted you to understand"}, {"start": 507.78000000000003, "end": 512.34, "text": " So the functional programming paradigm is that you can use the same array in place"}, {"start": 512.34, "end": 517.38, "text": " So the functional programming paradigm is something that Jax relies upon"}, {"start": 518.02, "end": 525.0600000000001, "text": " Finally, we have the fact that Jax handles random numbers completely differently compared to NumPy"}, {"start": 525.0600000000001, "end": 527.14, "text": " And there is a very very good reason for this"}, {"start": 527.14, "end": 528.82, "text": " We're going to see that in a moment"}, {"start": 528.82, "end": 536.5, "text": " But the key point here is you have to create this basically key which is a fancy name for like a state"}, {"start": 536.5, "end": 544.82, "text": " So whereas NumPy's pseudo random number generators are stateful, Jax is not stateful because it's"}, {"start": 545.38, "end": 549.38, "text": " As I said, that's a consequence of functional programming paradigm again"}, {"start": 549.38, "end": 553.14, "text": " So key is just a synonym for a state in Jax"}, {"start": 553.14, "end": 556.42, "text": " And what you have to do when you're generating random numbers"}, {"start": 556.42, "end": 561.78, "text": " You have to pass the state explicitly and not implicitly such like in NumPy"}, {"start": 561.78, "end": 564.9, "text": " So here we created 10 like random numbers"}, {"start": 564.9, "end": 566.8199999999999, "text": " We sampled them from a normal distribution"}, {"start": 567.62, "end": 571.62, "text": " And we can see here that we get 10 random numbers"}, {"start": 571.62, "end": 577.54, "text": " The thing I want you to notice here is that the type of this x array is of device array"}, {"start": 577.54, "end": 579.54, "text": " So that basically that's a trick"}, {"start": 579.54, "end": 583.22, "text": " And that's why Jax code is accelerator agnostic"}, {"start": 583.9399999999999, "end": 588.5799999999999, "text": " Basically, Jax automatically puts this array onto the accelerator"}, {"start": 588.5799999999999, "end": 593.78, "text": " So in this case, I'm using here my accelerator is set up to GPU"}, {"start": 593.78, "end": 598.98, "text": " If I were to set it to TPU, the code would be set directly to the TPU unit"}, {"start": 598.98, "end": 600.02, "text": " And that's very cool"}, {"start": 600.5799999999999, "end": 602.5799999999999, "text": " We're going to see a consequence of that a bit later"}, {"start": 602.5799999999999, "end": 607.3, "text": " Basically, you don't have to do the laborious to device from PyTorch syntax, etc"}, {"start": 607.3, "end": 609.6999999999999, "text": " So this story naturally leads me to fact number four"}, {"start": 609.6999999999999, "end": 614.02, "text": " And that's that Jax's accelerator agnostic same code runs everywhere"}, {"start": 614.02, "end": 618.98, "text": " You do not have to write particular syntax for a particular accelerators"}, {"start": 618.98, "end": 619.54, "text": " That's very cool"}, {"start": 619.54, "end": 625.4599999999999, "text": " Especially in the future where we see more and more companies being like building custom AI"}, {"start": 625.4599999999999, "end": 627.9399999999999, "text": " Accelerators and chips, not just TPUs and GPUs"}, {"start": 627.9399999999999, "end": 631.38, "text": " But like companies such as Graphcore, Cerebras, etc"}, {"start": 631.38, "end": 632.18, "text": " Let's see it in practice"}, {"start": 632.18, "end": 634.5, "text": " So here we define a Jax array"}, {"start": 635.14, "end": 637.4599999999999, "text": " Again, we see the syntax where we have the key"}, {"start": 637.4599999999999, "end": 642.18, "text": " We have to pass explicitly the key or the state of the pseudo random number generator"}, {"start": 642.18, "end": 646.3399999999999, "text": " And as you can notice, there is some difference in the API design here"}, {"start": 646.3399999999999, "end": 647.86, "text": " We have D type here"}, {"start": 647.86, "end": 649.86, "text": " We're using for NumPy, we're using S type"}, {"start": 649.86, "end": 651.0600000000001, "text": " There are some minor differences"}, {"start": 651.0600000000001, "end": 657.78, "text": " But like in general, syntaxes are very, very similar between Jax's NumPy and NumPy API"}, {"start": 658.98, "end": 662.1800000000001, "text": " I mentioned that we don't have to use two device anymore"}, {"start": 662.1800000000001, "end": 663.46, "text": " So that's in PyTorch"}, {"start": 663.46, "end": 667.86, "text": " So this array here is directly pushed to the GPU unit"}, {"start": 667.86, "end": 669.38, "text": " Whereas this one is on the CPU"}, {"start": 670.1, "end": 674.66, "text": " If we do some time profiling here, we can see that if let me run this thing"}, {"start": 674.66, "end": 679.86, "text": " Let me run this thing, basically this first line, line number 10"}, {"start": 680.74, "end": 685.06, "text": " We do like a basically matrix multiply directly on the GPU"}, {"start": 685.62, "end": 686.9, "text": " And that's very fast"}, {"start": 686.9, "end": 689.62, "text": " We can see that we're going to see results in a second"}, {"start": 689.62, "end": 693.06, "text": " The second line here, because we have a NumPy dot"}, {"start": 693.06, "end": 695.38, "text": " And because we have NumPy arrays which are on CPU"}, {"start": 696.02, "end": 697.3, "text": " It's going to be way slower"}, {"start": 697.3, "end": 702.26, "text": " And NumPy only works, the reason why Jax exists is because NumPy only works with CPUs"}, {"start": 702.26, "end": 705.06, "text": " And we do need to leverage existing accelerators"}, {"start": 705.7, "end": 711.86, "text": " The third line here does something like a modified version of the first two lines"}, {"start": 711.86, "end": 714.34, "text": " Basically, we're using Jax's dot product"}, {"start": 714.34, "end": 720.98, "text": " And we are using NumPy arrays which will cause this line to push the devices to the GPU"}, {"start": 720.98, "end": 723.22, "text": " And then do the dot product"}, {"start": 723.22, "end": 726.1, "text": " And because of that overhead, we're going to have a slower result"}, {"start": 726.1, "end": 727.62, "text": " Hopefully for this third line"}, {"start": 728.5, "end": 730.34, "text": " And let's see the numbers here"}, {"start": 730.34, "end": 732.74, "text": " So we have 27 milliseconds for the first line"}, {"start": 732.74, "end": 735.5400000000001, "text": " We have 437 milliseconds for the second line"}, {"start": 735.5400000000001, "end": 739.22, "text": " Because as I said, NumPy is very slow, it's running on CPU"}, {"start": 739.22, "end": 743.22, "text": " Finally, we have 97 milliseconds for this line here"}, {"start": 743.22, "end": 748.34, "text": " Because we have the overhead of sending the data from the CPU, from the host to the GPU unit"}, {"start": 749.38, "end": 752.26, "text": " And the final number 26 is the same as this one"}, {"start": 752.26, "end": 753.46, "text": " And that's this one here"}, {"start": 753.46, "end": 758.1, "text": " So basically what we do, we can explicitly, using this device put function"}, {"start": 758.1, "end": 761.86, "text": " We can push the DRA to the GPU"}, {"start": 761.86, "end": 766.74, "text": " And then now we have equivalent lines between this one and line number 10"}, {"start": 767.46, "end": 769.46, "text": " And that's it, that's why we have the same results here"}, {"start": 770.34, "end": 771.78, "text": " A couple of notes here"}, {"start": 771.78, "end": 775.38, "text": " First, I'm using GPU as a synonym for AI accelerator"}, {"start": 775.38, "end": 779.62, "text": " In reality, as I said, Jax's accelerator agnostic"}, {"start": 779.62, "end": 785.38, "text": " Which means basically, depending on what I set here in the runtime setting here"}, {"start": 785.38, "end": 789.9399999999999, "text": " We are going to run this either on GPU or TPU or whatever they have supported"}, {"start": 789.9399999999999, "end": 792.42, "text": " A second note is this block until ready"}, {"start": 792.42, "end": 797.9399999999999, "text": " You can notice that every time I'm using Jax's dot function, I'm using this block until ready"}, {"start": 797.9399999999999, "end": 803.3, "text": " The reason is Jax is using something called asynchronous dispatch system in the background"}, {"start": 803.3, "end": 806.18, "text": " Which means if I were to run this thing here"}, {"start": 806.18, "end": 807.9399999999999, "text": " So let me just copy paste this"}, {"start": 807.9399999999999, "end": 809.78, "text": " So if I were to copy paste this here"}, {"start": 809.78, "end": 815.4599999999999, "text": " If I were to run this and maybe like let me assign that to a variable"}, {"start": 815.4599999999999, "end": 818.42, "text": " And now I do some coding here, blah, blah, blah"}, {"start": 818.42, "end": 821.4599999999999, "text": " The thing is this line is going to immediately return"}, {"start": 822.18, "end": 825.06, "text": " And this code here is going to start executing"}, {"start": 825.06, "end": 828.98, "text": " Because this line actually just delegates the task to the accelerator"}, {"start": 828.98, "end": 832.42, "text": " And we're not blocked, we're not waiting for this line to execute"}, {"start": 832.42, "end": 839.14, "text": " Which means the variable, this temporary variable is still not filled in"}, {"start": 839.14, "end": 840.58, "text": " Let's put it that way"}, {"start": 840.58, "end": 843.9399999999999, "text": " That's a fairly neat feature that makes things even faster"}, {"start": 843.9399999999999, "end": 845.62, "text": " But you need to be cognizant of it"}, {"start": 845.62, "end": 847.86, "text": " Especially when you're doing this time profiling"}, {"start": 847.86, "end": 854.58, "text": " You want to measure the actual computation and not the time that's necessary to dispatch"}, {"start": 854.58, "end": 856.8199999999999, "text": " Just delegate the task to the accelerator"}, {"start": 856.8199999999999, "end": 857.78, "text": " Just be aware of that"}, {"start": 857.78, "end": 859.22, "text": " Nice, those were some basics"}, {"start": 859.22, "end": 861.9399999999999, "text": " Now we're going to cover the transform functions"}, {"start": 861.9399999999999, "end": 864.58, "text": " But before that, let me do a quick overview"}, {"start": 864.58, "end": 867.3, "text": " So Jax's AI accelerator agnostic"}, {"start": 867.3, "end": 869.54, "text": " Jax handles random numbers a bit differently"}, {"start": 869.54, "end": 872.42, "text": " We're going to see why that is important a bit later"}, {"start": 872.42, "end": 877.2199999999999, "text": " Basically reproducibility in the environments where you have multiple accelerators"}, {"start": 877.2199999999999, "end": 882.02, "text": " Where you're doing parallel programming is like possible with Jax"}, {"start": 882.02, "end": 886.9799999999999, "text": " Whereas NumPy has some problems there because it was designed mainly for CPU programs"}, {"start": 887.78, "end": 892.5799999999999, "text": " Then as I said, functional programming paradigm is something you need to get comfortable with"}, {"start": 892.5799999999999, "end": 896.66, "text": " And finally, basically the syntax is very similar to NumPy"}, {"start": 896.66, "end": 899.78, "text": " That's a high level overview of what Jax is"}, {"start": 899.78, "end": 904.42, "text": " And now we're going to get deeper into these transform functions"}, {"start": 904.42, "end": 907.6999999999999, "text": " So as I said, JIT compiles your functions using this XLA"}, {"start": 907.6999999999999, "end": 914.5799999999999, "text": " And basically doing that, it catches some functions and makes it very, very fast to execute our programs"}, {"start": 914.5799999999999, "end": 918.9, "text": " I'm going to just run this simple visualization function here"}, {"start": 918.9, "end": 923.86, "text": " And I'm going to show you a simple example where how we can JIT a function and make it faster"}, {"start": 923.86, "end": 926.74, "text": " So here I just defined this function"}, {"start": 926.74, "end": 929.7, "text": " It's a cellU function, it's just an activation function"}, {"start": 929.7, "end": 931.38, "text": " The details here are not that important"}, {"start": 931.38, "end": 934.58, "text": " That's why I have this visualize function here"}, {"start": 934.58, "end": 935.7, "text": " I'm just going to run it"}, {"start": 935.7, "end": 937.62, "text": " We're going to see how it looks like"}, {"start": 937.62, "end": 942.42, "text": " You can see here up until zero and that's defined here"}, {"start": 942.42, "end": 945.3000000000001, "text": " When X is greater than zero, we just have X"}, {"start": 945.3000000000001, "end": 947.86, "text": " Which means we have a linear function here"}, {"start": 947.86, "end": 953.3000000000001, "text": " And when the number X is smaller than zero, we just have X"}, {"start": 953.3, "end": 958.3399999999999, "text": " And smaller than zero, we have some combination of exponents here with some coefficients, etc"}, {"start": 958.3399999999999, "end": 959.38, "text": " But that's not important"}, {"start": 959.38, "end": 962.26, "text": " The important part here is we can transform cellU"}, {"start": 962.26, "end": 966.66, "text": " So we pass the function inside of this JIT transform function"}, {"start": 966.66, "end": 968.18, "text": " And we get a function back"}, {"start": 968.18, "end": 969.6999999999999, "text": " And this function is compiled"}, {"start": 969.6999999999999, "end": 971.4599999999999, "text": " Actually, it's not exactly compiled"}, {"start": 971.4599999999999, "end": 973.62, "text": " We first need to trace it by calling it once"}, {"start": 973.62, "end": 978.5, "text": " But for the sake of argument, we just have a super fast function right now here"}, {"start": 978.5, "end": 979.8599999999999, "text": " And let's benchmark it"}, {"start": 979.8599999999999, "end": 983.2199999999999, "text": " So let me allocate a vector of million data points"}, {"start": 983.22, "end": 986.4200000000001, "text": " We have million random numbers here"}, {"start": 986.4200000000001, "end": 989.14, "text": " And we pass them through the cellU function"}, {"start": 989.14, "end": 995.7, "text": " And we basically do a time profiling of both the normal version as well as the JIT version"}, {"start": 995.7, "end": 998.5, "text": " Let me run it again and let's see the numbers"}, {"start": 998.5, "end": 1002.9, "text": " So the thing here is this function is not overly complex"}, {"start": 1002.9, "end": 1004.9, "text": " So the results we get with JIT"}, {"start": 1004.9, "end": 1008.98, "text": " So we will not get a huge performance boost here"}, {"start": 1008.98, "end": 1011.5400000000001, "text": " But imagine once you start training neural networks"}, {"start": 1011.54, "end": 1018.26, "text": " Where JIT starts shining because then the optimizer, the compiler has a bigger freedom"}, {"start": 1018.26, "end": 1022.74, "text": " To do various, various optimization tricks such as fusing the operators"}, {"start": 1022.74, "end": 1027.46, "text": " Such as avoiding allocating certain temporary structures, etc"}, {"start": 1027.46, "end": 1031.22, "text": " And things get really, really fast like orders of magnitude faster"}, {"start": 1032.1, "end": 1038.1, "text": " Here we can see that basically we have 1.96 milliseconds using the normal version"}, {"start": 1038.1, "end": 1041.78, "text": " And then we have only 121 microseconds using the JIT version"}, {"start": 1041.78, "end": 1045.3, "text": " So even in this simple example, we get performance benefits"}, {"start": 1045.3, "end": 1048.26, "text": " So for now, I'm going to leave it here"}, {"start": 1048.26, "end": 1050.74, "text": " I'm going to stay like on a high level here"}, {"start": 1051.6999999999998, "end": 1053.86, "text": " We don't have to understand currently how JIT works"}, {"start": 1053.86, "end": 1057.86, "text": " But like you just need to know when you see JIT, JIT makes functions run fast"}, {"start": 1057.86, "end": 1059.1399999999999, "text": " Let me try to be this guy"}, {"start": 1062.58, "end": 1063.3799999999999, "text": " Awesome"}, {"start": 1063.3799999999999, "end": 1066.1799999999998, "text": " Okay, let's go ahead here"}, {"start": 1066.18, "end": 1068.74, "text": " Second transform function, very important, is grad"}, {"start": 1069.78, "end": 1074.3400000000001, "text": " Basically it does the magic of automatic differentiation for you"}, {"start": 1074.3400000000001, "end": 1078.9, "text": " So the same thing as dot backward in PyTorch if you're coming from the PyTorch world"}, {"start": 1079.78, "end": 1086.3400000000001, "text": " And like a short note here, differentiation can be manual"}, {"start": 1086.3400000000001, "end": 1090.74, "text": " So basically you have a function, you manually calculate on the paper"}, {"start": 1090.74, "end": 1094.98, "text": " How the derivatives look like"}, {"start": 1094.98, "end": 1098.26, "text": " And then you can code that knowledge into your function"}, {"start": 1098.26, "end": 1100.1, "text": " And that's a manual differentiation"}, {"start": 1100.1, "end": 1101.6200000000001, "text": " Symbolic one is very similar"}, {"start": 1101.6200000000001, "end": 1105.46, "text": " It's just an automatic way to automate the manual differentiation"}, {"start": 1105.46, "end": 1109.3, "text": " Whereby the program is using those rules such as product rule, etc"}, {"start": 1109.3, "end": 1111.94, "text": " To build up derivatives of a function"}, {"start": 1112.66, "end": 1115.46, "text": " And numeric functions are things like finite derivatives"}, {"start": 1115.46, "end": 1118.9, "text": " Where you use basically numeric methods to compute the derivatives"}, {"start": 1118.9, "end": 1121.46, "text": " And finally, the automatic one is the one we all love"}, {"start": 1121.46, "end": 1125.06, "text": " Which is used in every single deep learning framework we know of"}, {"start": 1125.94, "end": 1129.06, "text": " Okay, let's see how grad works"}, {"start": 1129.06, "end": 1131.46, "text": " So let's define a simple function here"}, {"start": 1131.46, "end": 1134.58, "text": " So it's just a sum of, as you can see here, logistic functions"}, {"start": 1134.58, "end": 1137.6200000000001, "text": " So this here inside of this is a logistic function or sigmoid"}, {"start": 1137.6200000000001, "end": 1143.38, "text": " And I just input, I just create like array of three values here, 0, 1, 2"}, {"start": 1143.38, "end": 1149.8600000000001, "text": " And I rename the function to loss because this could be used as a loss function"}, {"start": 1149.86, "end": 1151.6999999999998, "text": " Just giving some semantics"}, {"start": 1151.6999999999998, "end": 1156.02, "text": " And finally, we can do grad by just wrapping loss into grad"}, {"start": 1156.02, "end": 1158.4199999999998, "text": " Again, we are passing function into a function"}, {"start": 1158.4199999999998, "end": 1162.5, "text": " Which is something you see often in the functional program in paradigm"}, {"start": 1163.2199999999998, "end": 1165.06, "text": " That gives us the grad loss"}, {"start": 1165.06, "end": 1167.86, "text": " So the gradient of the loss function which we can evaluate"}, {"start": 1167.86, "end": 1171.62, "text": " So this is contrast this to PyTorch or TensorFlow"}, {"start": 1171.62, "end": 1173.1399999999999, "text": " Where you just do backward"}, {"start": 1173.1399999999999, "end": 1175.06, "text": " Here you actually get the function back"}, {"start": 1175.06, "end": 1177.62, "text": " And you can evaluate it at particular points"}, {"start": 1177.62, "end": 1183.4599999999998, "text": " And by default, I said here the grad will take the derivative of the first parameter"}, {"start": 1183.4599999999998, "end": 1186.1799999999998, "text": " But here we only have one parameter so that doesn't matter"}, {"start": 1186.1799999999998, "end": 1190.4199999999998, "text": " And let me kind of run this and see what we got"}, {"start": 1191.1399999999999, "end": 1195.2199999999998, "text": " So we have some numbers which mean nothing to us because the function is fairly complicated"}, {"start": 1195.2199999999998, "end": 1199.9399999999998, "text": " If I were to change this to this sum of squares function"}, {"start": 1199.9399999999998, "end": 1205.3, "text": " So if I do something like this, we basically can get interpretable results"}, {"start": 1205.3, "end": 1207.78, "text": " So let's do a manual derivative of this function"}, {"start": 1207.78, "end": 1213.86, "text": " We're going to get so that's x1 squared plus x2 squared plus x3 squared"}, {"start": 1213.86, "end": 1215.86, "text": " Remember we are passing three numbers here"}, {"start": 1215.86, "end": 1224.6599999999999, "text": " So if we do derivative of that, that's gonna be x1 times 2 plus x2 times 2 plus x3 times 2"}, {"start": 1224.6599999999999, "end": 1230.26, "text": " And if I were to run this, we're going to get 024 which makes sense because we have 012"}, {"start": 1230.26, "end": 1231.3799999999999, "text": " Let me print this out"}, {"start": 1231.38, "end": 1238.66, "text": " So print x, so we have 012 and we get 024 because of this"}, {"start": 1238.66, "end": 1241.46, "text": " Grad basically does because all of these are bundled inside of x"}, {"start": 1242.18, "end": 1246.18, "text": " Grad basically does derivative of the function with respect to x1"}, {"start": 1246.8200000000002, "end": 1250.66, "text": " Which is 2x1 and then the same for the other variables"}, {"start": 1250.66, "end": 1253.5400000000002, "text": " Originally, remember we have something like this"}, {"start": 1253.5400000000002, "end": 1257.8600000000001, "text": " We have this thing plus x2 squared plus x3 squared"}, {"start": 1257.8600000000001, "end": 1259.3000000000002, "text": " That's the original function"}, {"start": 1259.3, "end": 1261.86, "text": " This here is the derivative, manual derivative"}, {"start": 1262.4199999999998, "end": 1264.8999999999999, "text": " Okay, that's a simple example"}, {"start": 1264.8999999999999, "end": 1271.3, "text": " Let's now go and make sure that this thing actually does what we expect it to do"}, {"start": 1272.02, "end": 1275.62, "text": " We can do that using this numeric differentiation method"}, {"start": 1275.62, "end": 1277.22, "text": " The finite differences method"}, {"start": 1277.22, "end": 1280.58, "text": " If I were to run this, we're going to get some results"}, {"start": 1280.58, "end": 1284.34, "text": " There are some small errors here which are a natural consequence of the fact that we're"}, {"start": 1284.34, "end": 1287.3, "text": " storing numbers using a finite number of bits"}, {"start": 1287.3, "end": 1289.3, "text": " So how this works is very simple"}, {"start": 1289.3, "end": 1294.6599999999999, "text": " We take, so remember x is this array here, 012"}, {"start": 1294.6599999999999, "end": 1297.3799999999999, "text": " Length of x thus will return us 3"}, {"start": 1297.3799999999999, "end": 1300.4199999999998, "text": " We get the i is just the identity matrix"}, {"start": 1300.4199999999998, "end": 1304.34, "text": " And so we are basically iterating and taking one hot vectors here"}, {"start": 1304.34, "end": 1309.78, "text": " And basically this is here is the, this is the formula of the actual derivative"}, {"start": 1309.78, "end": 1315.06, "text": " So you nudge the input vector along a certain dimension a little bit"}, {"start": 1315.06, "end": 1317.46, "text": " In the positive direction, in the negative one"}, {"start": 1317.46, "end": 1322.02, "text": " And then you divide that by the 2 epsilon because yeah, that's the intensity of the nudge"}, {"start": 1322.02, "end": 1329.46, "text": " And if you don't understand this, the best way to play with Colab is to just insert a new code line here"}, {"start": 1329.46, "end": 1331.86, "text": " We can see what this exactly does"}, {"start": 1331.86, "end": 1333.46, "text": " So let me just paste this here"}, {"start": 1333.46, "end": 1336.34, "text": " So we have this thing here"}, {"start": 1336.34, "end": 1338.5, "text": " So we have, as you can see here, identity matrix"}, {"start": 1339.3799999999999, "end": 1343.86, "text": " And so this for loop is going to take one vector as a value"}, {"start": 1343.86, "end": 1348.02, "text": " One vector at a time and pass it into this function here"}, {"start": 1348.02, "end": 1355.3799999999999, "text": " And so you can see that basically we're going to evaluate f at the following"}, {"start": 1356.58, "end": 1359.9399999999998, "text": " So at least in the first iteration we're going to evaluate it at 0"}, {"start": 1359.9399999999998, "end": 1362.1799999999998, "text": " And then we're going to add this small epsilon here"}, {"start": 1362.1799999999998, "end": 1365.62, "text": " So epsilon and then we're going to have 1, 2"}, {"start": 1365.62, "end": 1367.4599999999998, "text": " And we're going to evaluate the function here"}, {"start": 1368.74, "end": 1373.4599999999998, "text": " We are going to subtract from that x minus epsilon"}, {"start": 1373.46, "end": 1375.7, "text": " So sorry, so that's 0 minus epsilon"}, {"start": 1376.3400000000001, "end": 1379.38, "text": " This thing will be constant for the first iteration"}, {"start": 1379.38, "end": 1381.06, "text": " And then we divide this by 2 epsilon"}, {"start": 1381.06, "end": 1384.26, "text": " And you can see that this is the definition of derivative itself"}, {"start": 1384.26, "end": 1387.3, "text": " And we do that for alongside every dimension"}, {"start": 1387.3, "end": 1391.38, "text": " That means in the second loop we'll have epsilon here instead of here"}, {"start": 1392.5, "end": 1393.8600000000001, "text": " And so on and so forth"}, {"start": 1393.8600000000001, "end": 1397.38, "text": " So minus epsilon here and yeah, you get the point"}, {"start": 1397.94, "end": 1402.82, "text": " So now we are certain that grad-brooks as expected"}, {"start": 1402.82, "end": 1405.22, "text": " Now let's see some fun examples"}, {"start": 1406.02, "end": 1410.4199999999998, "text": " We're going to define a simple second order polynomial function here"}, {"start": 1411.1399999999999, "end": 1412.58, "text": " x squared plus x plus 4"}, {"start": 1413.3, "end": 1415.22, "text": " I'm going to run this to visualize it"}, {"start": 1415.9399999999998, "end": 1417.22, "text": " It's always nice to visualize stuff"}, {"start": 1417.22, "end": 1420.6599999999999, "text": " It's at least to me it makes stuff a bit more like"}, {"start": 1422.8999999999999, "end": 1425.9399999999998, "text": " It gives them this gut feeling and it's easier to understand what's going on"}, {"start": 1427.06, "end": 1430.5, "text": " We're going to do higher order derivatives here"}, {"start": 1430.5, "end": 1434.02, "text": " Just to show you how a powerful grad function actually is"}, {"start": 1434.02, "end": 1438.66, "text": " So you can do grad of f and then you can do grad of the grad of f"}, {"start": 1438.66, "end": 1443.14, "text": " And then you can do grad of and you can do that n times whatever the number"}, {"start": 1443.14, "end": 1446.5, "text": " Like whatever the whatever n is"}, {"start": 1446.5, "end": 1447.54, "text": " Or you can just do it like this"}, {"start": 1447.54, "end": 1450.58, "text": " You can do it like grad of grad of grad of f"}, {"start": 1450.58, "end": 1451.78, "text": " This will also work"}, {"start": 1451.78, "end": 1454.1, "text": " This is just shorter so that's why I did it like this"}, {"start": 1454.1, "end": 1457.46, "text": " If we were to manually do a derivative of this polynomial"}, {"start": 1457.46, "end": 1460.18, "text": " We'll get the first derivative will give us 2x plus 1"}, {"start": 1460.18, "end": 1463.94, "text": " The second derivative so derivative of this one will just leave us with 2"}, {"start": 1463.94, "end": 1466.26, "text": " And finally derivative of a constant is 0"}, {"start": 1466.26, "end": 1469.46, "text": " That's why we have the third derivative is equal 0"}, {"start": 1469.46, "end": 1474.1000000000001, "text": " So if we print all the values here we expect to see because x is 1"}, {"start": 1474.1000000000001, "end": 1478.26, "text": " We expect to see 3, 2 and 0, right?"}, {"start": 1478.26, "end": 1480.18, "text": " So here it is 3, 2 and 0"}, {"start": 1480.18, "end": 1485.14, "text": " The first value is just the value of the function at the input x which is 1"}, {"start": 1485.14, "end": 1488.98, "text": " Which is as you can see here approximately 6"}, {"start": 1488.98, "end": 1490.66, "text": " Exactly 6, okay?"}, {"start": 1491.94, "end": 1492.74, "text": " So far so good"}, {"start": 1492.74, "end": 1495.46, "text": " The cool thing about this is we are very close to the math"}, {"start": 1495.46, "end": 1498.66, "text": " So you can basically as you see formula in some paper"}, {"start": 1498.66, "end": 1502.42, "text": " You can implement it much easier in JAX compared to PyTorch or TensorFlow"}, {"start": 1503.06, "end": 1504.02, "text": " It's very powerful"}, {"start": 1504.02, "end": 1509.38, "text": " We're gonna see how powerful this whole autodiff package of JAX is a bit later"}, {"start": 1509.94, "end": 1513.38, "text": " And now I'm gonna do a simple modification here"}, {"start": 1513.38, "end": 1515.46, "text": " Let's assume we have two inputs"}, {"start": 1515.46, "end": 1518.34, "text": " So let's assume we have y here"}, {"start": 1518.34, "end": 1520.98, "text": " So we have plus y squared"}, {"start": 1520.98, "end": 1523.86, "text": " And I'm just going to modify this lambda function"}, {"start": 1524.4199999999998, "end": 1529.06, "text": " I'm gonna skip the evaluation because the evaluation will fail now"}, {"start": 1529.06, "end": 1532.82, "text": " And now I'm going to do the following"}, {"start": 1532.82, "end": 1536.4199999999998, "text": " If we just do the grads of the function"}, {"start": 1537.86, "end": 1541.78, "text": " By default do the derivative with respect to the first parameter"}, {"start": 1541.78, "end": 1544.8999999999999, "text": " That means we're going to get this thing here exactly"}, {"start": 1544.8999999999999, "end": 1547.78, "text": " So we should expect exactly the same numbers"}, {"start": 1547.78, "end": 1552.8999999999999, "text": " Except for the fact that f will be different when we evaluate it at x"}, {"start": 1552.8999999999999, "end": 1557.62, "text": " So this obviously will not work until I enter some numbers"}, {"start": 1557.62, "end": 1561.3799999999999, "text": " So this is going to be like let me say it's gonna be 1"}, {"start": 1561.3799999999999, "end": 1563.62, "text": " And then I'm gonna modify this here"}, {"start": 1563.62, "end": 1571.54, "text": " So 1 x y x y and finally x y here"}, {"start": 1572.26, "end": 1573.22, "text": " That should work"}, {"start": 1573.22, "end": 1574.5, "text": " Let's see if it works or not"}, {"start": 1575.54, "end": 1577.3, "text": " Yeah, we get the numbers as expected"}, {"start": 1577.3, "end": 1579.3, "text": " So we get the same results here"}, {"start": 1579.3, "end": 1582.74, "text": " And we get 7 because once you evaluate it at y equals 1"}, {"start": 1582.74, "end": 1584.82, "text": " This will be 6 plus 1 that's 7"}, {"start": 1586.18, "end": 1589.46, "text": " So how can we do derivative with respect to y?"}, {"start": 1589.46, "end": 1594.98, "text": " The only thing we need to do here is arc nums equals 1 instead"}, {"start": 1594.98, "end": 1596.6599999999999, "text": " I think this is the syntax"}, {"start": 1596.6599999999999, "end": 1597.46, "text": " Let me try it out"}, {"start": 1598.34, "end": 1601.22, "text": " Let me just ignore this for the time being"}, {"start": 1601.22, "end": 1602.82, "text": " I'm just gonna delete this"}, {"start": 1603.3799999999999, "end": 1605.3799999999999, "text": " And let me print this out"}, {"start": 1605.3799999999999, "end": 1606.8999999999999, "text": " So we get the number 2 here"}, {"start": 1606.9, "end": 1608.26, "text": " Which is correct result"}, {"start": 1608.26, "end": 1612.18, "text": " If we do a differentiation of this polynomial with respect to y"}, {"start": 1612.18, "end": 1614.9, "text": " We expect to get 2 times y"}, {"start": 1614.9, "end": 1618.02, "text": " Since y is equal 1 we get 2 here"}, {"start": 1618.02, "end": 1622.5, "text": " So the thing that confuses me here is I'm not sure why we are getting device array here"}, {"start": 1622.5, "end": 1625.7800000000002, "text": " Whereas previously we had pure floats on the CPU"}, {"start": 1626.5800000000002, "end": 1628.1000000000001, "text": " I'm not sure about that"}, {"start": 1628.1000000000001, "end": 1630.5, "text": " Let me check if I can do it like this"}, {"start": 1630.5, "end": 1636.9, "text": " Yep, for some reason once I add a comma here"}, {"start": 1636.9, "end": 1638.82, "text": " It's returning me device array"}, {"start": 1638.82, "end": 1639.7, "text": " Whatever"}, {"start": 1639.7, "end": 1642.1, "text": " Obviously this is now df dy"}, {"start": 1642.1, "end": 1643.94, "text": " But yeah, you get the point"}, {"start": 1644.98, "end": 1648.18, "text": " Now let me show you how powerful the autodef engine is"}, {"start": 1648.18, "end": 1654.1, "text": " So aside from grad which can do derivatives with respect for the scalar output functions"}, {"start": 1654.1, "end": 1660.02, "text": " Jacobians, so these jack forward and jack rev functions can find Jacobians"}, {"start": 1660.02, "end": 1664.82, "text": " So Jacobians basically can evaluate derivatives even for vector-valued functions"}, {"start": 1664.82, "end": 1666.02, "text": " That's the only difference"}, {"start": 1668.02, "end": 1669.6999999999998, "text": " That's pretty much the only difference"}, {"start": 1669.6999999999998, "end": 1673.2199999999998, "text": " So here let me take another function"}, {"start": 1673.2199999999998, "end": 1676.02, "text": " Basically we have a simple paraboloid"}, {"start": 1676.02, "end": 1679.54, "text": " If we google it here we can see how it looks like"}, {"start": 1679.54, "end": 1684.26, "text": " It looks like a cup 2D surface in a 3D space"}, {"start": 1684.98, "end": 1688.8999999999999, "text": " And basically if we manually calculate the derivatives"}, {"start": 1688.8999999999999, "end": 1692.74, "text": " So df dx will be 2x, df dy will be 2y"}, {"start": 1692.74, "end": 1695.54, "text": " So we expect Jacobian to look something like this"}, {"start": 1695.54, "end": 1699.54, "text": " If we continue on and figure out the second order derivatives"}, {"start": 1699.54, "end": 1700.82, "text": " We can get the Hessian number"}, {"start": 1700.82, "end": 1704.42, "text": " So we can get the second order derivative with respect to x is 2"}, {"start": 1704.42, "end": 1707.86, "text": " As you can see here for y it's also going to be 2"}, {"start": 1707.86, "end": 1712.5, "text": " And if we do dx and then dy because this does not depend on y anymore"}, {"start": 1712.5, "end": 1713.86, "text": " We're going to get zeros"}, {"start": 1713.86, "end": 1716.58, "text": " And finally how you form the Hessian matrix"}, {"start": 1716.58, "end": 1718.8999999999999, "text": " And you don't need to worry about what Hessian is"}, {"start": 1718.8999999999999, "end": 1721.1399999999999, "text": " It's basically just a collection"}, {"start": 1721.1399999999999, "end": 1728.5, "text": " It's a matrix of as you can see here various derivatives of the multivariate function"}, {"start": 1728.5, "end": 1732.1799999999998, "text": " You can simply define it using the jack rev and jack forward function"}, {"start": 1732.1799999999998, "end": 1733.2199999999998, "text": " We additionally did it"}, {"start": 1733.2199999999998, "end": 1736.74, "text": " You can see how this nicely composes and we can get the results"}, {"start": 1736.74, "end": 1742.98, "text": " So the reason they're using both forward and rev is because it's just optimization like detail"}, {"start": 1742.98, "end": 1747.22, "text": " Because one of these the rev one works with wide matrices"}, {"start": 1747.22, "end": 1750.98, "text": " Whereas the jack forward works with tall matrices"}, {"start": 1750.98, "end": 1755.78, "text": " I'm going to link a video down in the description which nicely explains why this is"}, {"start": 1755.78, "end": 1761.54, "text": " Anyways we get the numbers so that's 2, 2 for the Jacobian as we expected here"}, {"start": 1761.54, "end": 1765.06, "text": " So if we plug in the numbers so for the input 1, 1"}, {"start": 1765.06, "end": 1767.54, "text": " We can see that this evaluates to 2, 2"}, {"start": 1767.54, "end": 1770.8999999999999, "text": " And finally for the Hessian we get 2, 0, 0, 2"}, {"start": 1770.8999999999999, "end": 1774.74, "text": " So anyways I just wanted to show you that this is possible"}, {"start": 1774.74, "end": 1778.82, "text": " And if you want to dig into more details there documentation is very nice"}, {"start": 1778.82, "end": 1782.1, "text": " And yeah let's continue with the final example for grad"}, {"start": 1782.1, "end": 1785.46, "text": " I just took this edge case function so the absolute value of x"}, {"start": 1785.46, "end": 1787.94, "text": " And let's see how how how Jax handles it"}, {"start": 1787.94, "end": 1789.94, "text": " Basically this is how the function looks like"}, {"start": 1789.94, "end": 1791.46, "text": " I just visualized it here"}, {"start": 1791.46, "end": 1796.9, "text": " And I printed the values at minus 1 and 1 which would be 1 and 1"}, {"start": 1796.9, "end": 1798.9, "text": " And you can see the the results here"}, {"start": 1798.9, "end": 1802.9, "text": " And finally I printed the I found the gradient of this function"}, {"start": 1802.9, "end": 1805.7, "text": " And you can see it's not differentiable at this point at 0"}, {"start": 1805.7, "end": 1807.7, "text": " It's the gradient is undefined"}, {"start": 1807.7, "end": 1814.5, "text": " But like if we were to evaluate the function at minus 1 we get"}, {"start": 1814.5, "end": 1818.26, "text": " So as you can see here the derivative will be minus 1 for all the numbers here"}, {"start": 1818.26, "end": 1820.26, "text": " 1 for all numbers here"}, {"start": 1820.26, "end": 1822.26, "text": " The interesting point is actually 0"}, {"start": 1822.26, "end": 1826.26, "text": " Because as you can see this one returns the obvious result minus 1"}, {"start": 1826.26, "end": 1828.26, "text": " So this one here is kind of interesting"}, {"start": 1828.26, "end": 1836.26, "text": " If we were to delete this part and just kind of add a small number to 0"}, {"start": 1836.26, "end": 1840.26, "text": " We'll get 1 as the output"}, {"start": 1840.26, "end": 1846.26, "text": " If we were to do the same thing just on the negative side we're going to get minus 1"}, {"start": 1846.26, "end": 1853.78, "text": " So basically what they've done is they defined for 0 is going to be the value is going to be defined is going to be 1"}, {"start": 1853.78, "end": 1857.22, "text": " That's just a convention and that's how Jax deals with it"}, {"start": 1857.22, "end": 1862.26, "text": " We didn't get any exception so yeah I guess I guess that's good at least in some settings"}, {"start": 1862.26, "end": 1864.26, "text": " Finally let's jump to Vmap"}, {"start": 1864.26, "end": 1870.26, "text": " Basically when you see Vmap you need to think about it the main value prop of Vmap is the following"}, {"start": 1870.26, "end": 1873.86, "text": " Writer functions as if you were dealing with a single data point"}, {"start": 1873.86, "end": 1878.26, "text": " And this is going to become more clear as we go through these examples"}, {"start": 1878.26, "end": 1886.26, "text": " Let's say we have a matrix W which is basically weights of you can treat it as weights of a linear neural network layer"}, {"start": 1886.26, "end": 1890.26, "text": " We input some state we set the shape to 150, 100"}, {"start": 1890.26, "end": 1897.26, "text": " And we now create this batch X you can treat this as maybe 10 images a batch of 10 flattened images"}, {"start": 1897.26, "end": 1903.26, "text": " So it's 100 that means we had like a 10 times 10 pixel image it's 100 when you once you flattened out"}, {"start": 1903.26, "end": 1908.86, "text": " So now if we apply the matrix we do the dot product between W and X"}, {"start": 1908.86, "end": 1915.86, "text": " This is basically a simulating what the linear layer is doing in the background when it's processing the input data"}, {"start": 1915.86, "end": 1923.86, "text": " Now the trick here is this only transforms the single a single image it cannot handle the batch"}, {"start": 1923.86, "end": 1930.86, "text": " Otherwise it will crash because we're trying to multiply 150, 100 which is W"}, {"start": 1930.86, "end": 1937.4599999999998, "text": " And we're trying to multiply that with 10, 100 that's going to fail"}, {"start": 1937.4599999999998, "end": 1941.4599999999998, "text": " So this thing does not work for the batch so okay let me run it"}, {"start": 1941.4599999999998, "end": 1945.86, "text": " So how would we go about making this work for a batch of images?"}, {"start": 1945.86, "end": 1951.4599999999998, "text": " So that's because that's what we care about and because that's way more efficient as you may know"}, {"start": 1951.4599999999998, "end": 1954.86, "text": " Especially on GPUs and accelerators"}, {"start": 1954.86, "end": 1965.4599999999998, "text": " Basically the most naive approach would be to iterate through the images and then call this function that knows how to handle single data points"}, {"start": 1965.4599999999998, "end": 1971.06, "text": " And then we just stack the results and basically we can see that this will work"}, {"start": 1971.06, "end": 1979.6599999999999, "text": " And the thing is it's very very slow because you want to avoid doing for loops you want to vectorize your functions"}, {"start": 1979.6599999999999, "end": 1984.4599999999998, "text": " And that's why this runs 5.54 milliseconds"}, {"start": 1984.46, "end": 1989.06, "text": " Now let's see a bit more a better approach of doing this"}, {"start": 1989.06, "end": 1995.8600000000001, "text": " So this would be a better approach but as you can see we had to completely rewrite the function we had above"}, {"start": 1995.8600000000001, "end": 2000.46, "text": " So we had to swap as you can see here we had to swap W with X"}, {"start": 2000.46, "end": 2005.46, "text": " Now we have batch of X here and we have W here and we additionally had to do a transpose"}, {"start": 2005.46, "end": 2008.8600000000001, "text": " So now this will the shapes will match and everything will work fine"}, {"start": 2008.86, "end": 2016.6599999999999, "text": " So we have 100 and we try to multiply that with transpose version of W which is 100 150"}, {"start": 2016.6599999999999, "end": 2023.26, "text": " And that gives us as you can see 10 150 which is what we expected right?"}, {"start": 2023.26, "end": 2028.4599999999998, "text": " So everything is fine now the problem with this and it's additional jitted so it's going to be really fast"}, {"start": 2028.4599999999998, "end": 2035.86, "text": " Now the problem with this is as you can see we have to write depending on whether we handle singular cases or a batch of data"}, {"start": 2035.86, "end": 2040.4599999999998, "text": " We have to write completely different functions and that's not desirable"}, {"start": 2040.4599999999998, "end": 2048.06, "text": " Although you can see this is way faster. So it's only 103 microseconds. Whereas this is 5.54 milliseconds"}, {"start": 2048.06, "end": 2053.06, "text": " Now this is the result that Jacks offers us using something called V map"}, {"start": 2053.06, "end": 2057.2599999999998, "text": " You just take the function that handles a single data point"}, {"start": 2057.2599999999998, "end": 2061.2599999999998, "text": " You wrap it into this you transform it using this V map"}, {"start": 2061.26, "end": 2066.26, "text": " And you can now pass the batch of your data without any other modifications"}, {"start": 2066.26, "end": 2070.46, "text": " So that's that's very cool. If I run this"}, {"start": 2070.46, "end": 2074.46, "text": " Let's see how how fast this thing is. So it's around 100 microseconds"}, {"start": 2074.46, "end": 2077.6600000000003, "text": " I think the last time I ran this this was actually faster than this one"}, {"start": 2077.6600000000003, "end": 2080.6600000000003, "text": " But I guess there are some some variants inside of there"}, {"start": 2080.6600000000003, "end": 2083.6600000000003, "text": " Okay, so I went ahead and reran this thing again"}, {"start": 2083.6600000000003, "end": 2087.86, "text": " And now we can see that we have 180 microseconds for this function"}, {"start": 2087.86, "end": 2093.6600000000003, "text": " Whereas we have 143 microseconds for this one. So obviously there is some variance in here"}, {"start": 2093.6600000000003, "end": 2098.06, "text": " But the point is here you have a much simpler way to write these batched"}, {"start": 2098.86, "end": 2104.26, "text": " Functions and it's as as efficient as by like doing this laborious work"}, {"start": 2104.26, "end": 2111.86, "text": " So basically what what V map does in the background is it takes the the for loops and packs them into this Alex"}, {"start": 2112.26, "end": 2115.46, "text": " API mid-level layer of the of the API"}, {"start": 2115.46, "end": 2118.76, "text": " and yeah, so now let me let me go ahead and"}, {"start": 2119.76, "end": 2121.76, "text": " Modify this example a little bit"}, {"start": 2121.76, "end": 2124.0, "text": " So we are going to instead of passing just X"}, {"start": 2124.0, "end": 2131.36, "text": " Let's try and pass both the W as well as the X because that's more I guess in in line with the functional"}, {"start": 2131.96, "end": 2137.66, "text": " Programming paradigm. So let me rerun this cell and now let's modify this one to to accept"}, {"start": 2138.26, "end": 2143.56, "text": " like both the W as well as the as the batch of data so"}, {"start": 2143.56, "end": 2150.7599999999998, "text": " We'll need to add this thing here again. So if I were to run this right now, it will crash and we'll see why"}, {"start": 2151.36, "end": 2153.36, "text": " so it basically says"}, {"start": 2153.7599999999998, "end": 2155.7599999999998, "text": " It basically says that W"}, {"start": 2155.7599999999998, "end": 2162.7599999999998, "text": " We're trying to to tell to V map that W has a batch dimension where whereas it does not have it's just a it's just a matrix"}, {"start": 2162.7599999999998, "end": 2165.7599999999998, "text": " That's supposed to represent the linear layer"}, {"start": 2165.7599999999998, "end": 2169.16, "text": " So what we need to do is add the in axis argument here"}, {"start": 2169.16, "end": 2175.56, "text": " So in axis like this and then we say none because W does not have a batch dimension"}, {"start": 2175.56, "end": 2181.3599999999997, "text": " And we need to specify the batch dimension of the second input and that's zero. So now this should work"}, {"start": 2181.3599999999997, "end": 2183.3599999999997, "text": " Let me let me try and rerun it again"}, {"start": 2185.3599999999997, "end": 2187.3599999999997, "text": " So fingers crossed"}, {"start": 2187.3599999999997, "end": 2189.7599999999998, "text": " And yeah, it did work. So basically that's it"}, {"start": 2189.7599999999998, "end": 2194.96, "text": " Now you saw in a bit more detail how V map works and this is the first part of this video"}, {"start": 2194.96, "end": 2197.16, "text": " So we basically saw the basics of Jax"}, {"start": 2197.16, "end": 2204.16, "text": " We saw how to use the the main transform functions such as JIT such as grad and finally V map and now let's dig"}, {"start": 2204.16, "end": 2210.16, "text": " Even deeper and understand the intricacies of how JIT works because that will help you enormously"}, {"start": 2211.16, "end": 2214.16, "text": " Debugging these Jax programs. So let's let's continue here. So"}, {"start": 2215.16, "end": 2222.16, "text": " We basically have as I said Jax has this onion like API layer API structure and that's that's I guess pretty much"}, {"start": 2222.16, "end": 2224.16, "text": " Always the case but still"}, {"start": 2224.16, "end": 2229.16, "text": " Anyways, we have numpy as the as the highest level then we have lax as the mid level and finally the XLA"}, {"start": 2230.16, "end": 2235.16, "text": " So lax API is stricter and more powerful and it's a simple Python wrapper around XLA"}, {"start": 2235.16, "end": 2237.16, "text": " So let's see what it means when I say stricter"}, {"start": 2237.16, "end": 2242.16, "text": " So if we were to add in the numpy level of the API"}, {"start": 2242.16, "end": 2244.16, "text": " This thing can be tolerated"}, {"start": 2244.16, "end": 2249.16, "text": " So one plus one like one as the integer plus one as the integer"}, {"start": 2249.16, "end": 2256.16, "text": " Like one as the integer plus one as the float will work. But once we get to the mid level"}, {"start": 2257.16, "end": 2262.16, "text": " Here we need to be explicit about the types. Otherwise we'll have we'll get an error"}, {"start": 2262.16, "end": 2266.16, "text": " So if I run this we'll see that this is printed"}, {"start": 2266.16, "end": 2272.16, "text": " So one plus one is okay, but here add requires arguments to have the same D types got int 32 and float 32"}, {"start": 2272.16, "end": 2277.16, "text": " The reason they've done this done it this way is to be more error robust"}, {"start": 2277.16, "end": 2283.16, "text": " Finally, we have the fact that the lax layer is obviously more powerful"}, {"start": 2283.16, "end": 2289.16, "text": " Although as a trade-off, it's it's less user friendly, which is kind of obvious as an example here. We have X and Y"}, {"start": 2289.16, "end": 2297.16, "text": " We have basically we want to do a convolution a 1d convolution between the signals X and Y and in numpy"}, {"start": 2297.16, "end": 2300.16, "text": " API this would be done like this"}, {"start": 2300.16, "end": 2303.16, "text": " We just call the convol function and we get the result"}, {"start": 2303.16, "end": 2308.16, "text": " On the other hand, once you get into the lax land, you have to use this conv general dilated"}, {"start": 2308.16, "end": 2312.16, "text": " Which is way more powerful as you can see has much more options"}, {"start": 2312.16, "end": 2317.16, "text": " 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"}, {"start": 2317.16, "end": 2320.16, "text": " You have to be explicit about the the types etc"}, {"start": 2320.16, "end": 2325.16, "text": " And if I run this we can see the results here as well as here will be the same"}, {"start": 2325.16, "end": 2330.16, "text": " Like this this assert should make sure that that we get the same results"}, {"start": 2330.16, "end": 2333.16, "text": " And as you can see it works perfectly now"}, {"start": 2333.16, "end": 2340.16, "text": " Final remark here is you can see here that this result returns batched results"}, {"start": 2340.16, "end": 2344.16, "text": " So that's why we have to index 0 0 to get the actual result"}, {"start": 2344.16, "end": 2348.16, "text": " As I said, lax API is just a thin wrapper around thin"}, {"start": 2348.16, "end": 2355.16, "text": " I mean, it's a wrapper around XLA so you can actually find the the XLA function that's going to be eventually called here"}, {"start": 2355.16, "end": 2358.16, "text": " And that's this conv the general padding"}, {"start": 2358.16, "end": 2361.16, "text": " It's in C++ you can see the arguments etc"}, {"start": 2361.16, "end": 2368.16, "text": " So if you ever need to dig a bit deeper and optimize something really really really hard"}, {"start": 2368.16, "end": 2372.16, "text": " Then you TensorFlow documentation here got you covered"}, {"start": 2372.16, "end": 2375.16, "text": " So let's get back here and continue"}, {"start": 2375.16, "end": 2378.16, "text": " That was the short mention of the API"}, {"start": 2378.16, "end": 2380.16, "text": " Now let's finally understand how it works"}, {"start": 2380.16, "end": 2384.16, "text": " So again, I won't get into more details here"}, {"start": 2384.16, "end": 2392.16, "text": " The whole point of this is to show that JIT functions are faster compared to non-JIT versions of the function"}, {"start": 2392.16, "end": 2396.16, "text": " Here we just normalize the matrix X like column wise"}, {"start": 2396.16, "end": 2405.16, "text": " 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"}, {"start": 2405.16, "end": 2410.16, "text": " And finally, the difference here is again not that big because this is a super simple function"}, {"start": 2410.16, "end": 2417.16, "text": " The more complex the function is the bigger the differences here will be between the JIT version and between the normal version"}, {"start": 2417.16, "end": 2425.16, "text": " So here we have the JIT version and again we're using a block until ready because of the asynchronous dispatch and we see the results"}, {"start": 2425.16, "end": 2428.16, "text": " Okay, so let's see"}, {"start": 2428.16, "end": 2432.16, "text": " In order to understand JIT it may be useful to understand when it fails"}, {"start": 2432.16, "end": 2441.16, "text": " Which functions cannot, so which class of functions we cannot JIT will tell you, will help you understand how JIT actually works behind the curtains"}, {"start": 2441.16, "end": 2444.16, "text": " So if I were to run this thing it's gonna crash"}, {"start": 2444.16, "end": 2455.16, "text": " So we are creating a simple vector of 10 random elements and we pass that vector into this get negatives and we call it"}, {"start": 2455.16, "end": 2461.16, "text": " So now because we don't have any JIT we're gonna get a result back and everything works as expected"}, {"start": 2461.16, "end": 2468.16, "text": " But if I were to call the JIT version of the function by wrapping this into this JIT transform"}, {"start": 2468.16, "end": 2477.16, "text": " If we run this we're going to get like error and it says here array boolean indices must be concrete, got shaped array, blah blah blah"}, {"start": 2477.16, "end": 2484.16, "text": " So basically what happens here is that depending on the content of X, so depending on the values of X"}, {"start": 2484.16, "end": 2491.16, "text": " We are going to get an output that's going to vary in its shape and that's something that's not tolerable inside the JIT world"}, {"start": 2491.16, "end": 2496.16, "text": " Let's slowly understand why this is the case, so let's start with this function"}, {"start": 2496.16, "end": 2506.16, "text": " We have a simple function f which just does like a dot product between X and Y and returns result and in the meanwhile it also prints some intermediate variables here"}, {"start": 2506.16, "end": 2514.16, "text": " So I prepare the variable X and Y, just random vectors and matrix here and we call the function here"}, {"start": 2514.16, "end": 2520.16, "text": " And then we again call, we have a second call here and let's see what happens after the second call"}, {"start": 2520.16, "end": 2527.16, "text": " So I'm gonna run this, there's a couple things happening here, first thing first is the first time you run a JIT function"}, {"start": 2527.16, "end": 2532.16, "text": " So that's line 14, what JIT does in the background is something called tracing"}, {"start": 2532.16, "end": 2542.16, "text": " So it basically takes, instead of inputting the actual values of X and Y, it creates this abstract like tracer value"}, {"start": 2542.16, "end": 2549.16, "text": " Let's call it that way, so it's basically like a placeholder variable that has a specified shape and specified data type"}, {"start": 2549.16, "end": 2556.16, "text": " So you can see here basically once we print X, it outputs here instead of the value of X, it outputs traced shaped array"}, {"start": 2556.16, "end": 2566.16, "text": " It's basically float 32 because we had random numbers sampled from Gaussian, we have the shape here 3, 4 and we have 4 for Y"}, {"start": 2566.16, "end": 2572.16, "text": " So basically you can see these are actually passed into the function the first time JIT is called"}, {"start": 2572.16, "end": 2582.16, "text": " And using that JIT can understand how the shapes are morphing going through the function and finally what the output shape of the function is"}, {"start": 2582.16, "end": 2588.16, "text": " And you can see here the output is also float 32 and the shape is 3"}, {"start": 2588.16, "end": 2598.16, "text": " We finally get the results here, the second time you call the function something funny happens and that has to do with the functional programming paradigm"}, {"start": 2598.16, "end": 2608.16, "text": " Basically because print functions are a side effect because we are returning the result from this function by other means than through the actual output here"}, {"start": 2608.16, "end": 2615.16, "text": " So we are printing and that's a side effect and because of that JIT is just going to ignore all of those side effects"}, {"start": 2615.16, "end": 2622.16, "text": " And the second time you call it because it's going to call the compile function, we won't see any printing"}, {"start": 2622.16, "end": 2629.16, "text": " And the whole point is this caching mechanism I just explained, so that's why JIT functions are so fast"}, {"start": 2629.16, "end": 2634.16, "text": " So I mentioned here anytime we get the same shapes and types we just call the compile function"}, {"start": 2634.16, "end": 2644.16, "text": " So that means if I were to now call whatever, so basically for the whole class of inputs that have this shape and this data type"}, {"start": 2644.16, "end": 2647.16, "text": " And no matter the value we are always going to call the compile function"}, {"start": 2647.16, "end": 2653.16, "text": " But if I were to call this function again but this time with a different shape, so let's do something like this"}, {"start": 2653.16, "end": 2663.16, "text": " So x3 maybe 3, 5 and this will be 5 and let me just map this like this and I'm going to call the function again and print the results"}, {"start": 2663.16, "end": 2671.16, "text": " So x3, y3, so what will happen here is that this time we'll again trigger the compilation because the shape changed here"}, {"start": 2671.16, "end": 2678.16, "text": " And so JIT is smart enough to retrace it and as you can see here we have the tracing happen here again"}, {"start": 2678.16, "end": 2685.16, "text": " Hopefully you get a better picture of how JIT now works but we are going to continue and understand this in even more depth"}, {"start": 2685.16, "end": 2691.16, "text": " So now we have the same function as above, we just omit the print functions which are the side effects"}, {"start": 2691.16, "end": 2698.16, "text": " And this is how you should write your functions when you want to use them with JAKS transform functions such as JIT"}, {"start": 2698.16, "end": 2700.16, "text": " So no side effects"}, {"start": 2700.16, "end": 2706.16, "text": " If I were to print this and we are going to have something called JAKSper here, so that's JAKS expression"}, {"start": 2706.16, "end": 2714.16, "text": " And that's basically the, let's call it like a flow model that JIT creates in the background when it does its tracing procedure"}, {"start": 2714.16, "end": 2718.16, "text": " So if I were to run this, let's see what we have here"}, {"start": 2718.16, "end": 2727.16, "text": " So you can see what happens, so it creates this abstract grammar and you basically have C, so add 1 to A"}, {"start": 2727.16, "end": 2734.16, "text": " So that's going to be, and A is basically the first argument, so that's X, so it's basically creating a placeholder for this thing here"}, {"start": 2734.16, "end": 2739.16, "text": " Then it calls D B plus 1, so that's Y plus 1, this is going to be B"}, {"start": 2739.16, "end": 2746.16, "text": " And then it calls the general dot product, so that's this thing, with C and D, as you can see here, C and D"}, {"start": 2746.16, "end": 2750.16, "text": " And yeah, and it returns back the E, which is this result here"}, {"start": 2750.16, "end": 2756.16, "text": " So you can go into docs and understand in more detail every single part of this syntax of this grammar"}, {"start": 2756.16, "end": 2764.16, "text": " But like, now you understand a bit better how JIT, when it does the tracing, creates this type of a flow in the background"}, {"start": 2764.16, "end": 2768.16, "text": " And this can be compiled using XLA, okay?"}, {"start": 2768.16, "end": 2771.16, "text": " Let's see another example of a failure"}, {"start": 2771.16, "end": 2777.16, "text": " This time, what we do is we pass this argument neck and we condition upon it"}, {"start": 2777.16, "end": 2784.16, "text": " And depending on the value, whether it's true or false, we're going to return minus X or plus X"}, {"start": 2784.16, "end": 2796.16, "text": " And remember, the first time we call a JIT function, it's going to input abstract shape and like a data type"}, {"start": 2796.16, "end": 2801.16, "text": " It won't have information about the value, and that's why this thing is going to fail"}, {"start": 2801.16, "end": 2808.16, "text": " So it says here, blah blah blah, abstract tracer, value encountered, where concrete value is expected"}, {"start": 2808.16, "end": 2812.16, "text": " The problem arose in the bool function, so that's basically what happens"}, {"start": 2812.16, "end": 2817.16, "text": " And in order to avoid this, what we can do is use these static arguments"}, {"start": 2817.16, "end": 2823.16, "text": " By making an argument static, so we say here, hey, first argument, so that's neck, is going to be static"}, {"start": 2823.16, "end": 2831.16, "text": " What we do by doing this is once the JIT does the tracing procedure, it will not use this abstract tracer object"}, {"start": 2831.16, "end": 2833.16, "text": " It's going to use the actual value"}, {"start": 2833.16, "end": 2837.16, "text": " So we are kind of lowering the level of abstraction here while doing the tracing"}, {"start": 2837.16, "end": 2846.16, "text": " Which is going to kind of constrain the class of inputs where this compile function can be called from the cache"}, {"start": 2846.16, "end": 2851.16, "text": " But in return, yeah, in return, we get it to work, I guess"}, {"start": 2851.16, "end": 2855.16, "text": " Okay, so let's call this thing, so this should now work"}, {"start": 2855.16, "end": 2862.16, "text": " And the thing I want you to understand here again is that the first time we call it with true here, we do the tracing"}, {"start": 2862.16, "end": 2869.16, "text": " And then the second time we call it with true, because nothing changed, basically again we don't have the tracing procedure"}, {"start": 2869.16, "end": 2876.16, "text": " Once we switch this to false, we again trigger the tracing, so now we can call this function for any integer 32"}, {"start": 2876.16, "end": 2882.16, "text": " And we'll have two cached functions, so that's like, which run really fast"}, {"start": 2882.16, "end": 2885.16, "text": " Let's continue analyzing the failures, the third failure is here"}, {"start": 2885.16, "end": 2894.16, "text": " If I were to run this, let's see what we get, blah blah blah shapes must be 1D sequences of concrete values of integer types"}, {"start": 2894.16, "end": 2898.16, "text": " Okay, so what happened here is, let me try and print some values again"}, {"start": 2898.16, "end": 2909.16, "text": " Print x, let's print x shape, let's print this whole product thing"}, {"start": 2909.16, "end": 2913.16, "text": " And I think that's going to work, so let's try and print that"}, {"start": 2913.16, "end": 2922.16, "text": " 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"}, {"start": 2922.16, "end": 2931.16, "text": " And then this thing here is again a trace, so this trace object, and we're trying to pass the trace object into reshape, which expects a concrete value"}, {"start": 2931.16, "end": 2933.16, "text": " That's why this thing is crashing again"}, {"start": 2933.16, "end": 2942.16, "text": " Let's see what the solution is, we can basically use NumPy prod instead of Jack's product"}, {"start": 2942.16, "end": 2949.16, "text": " So this may be a little bit confusing, basically as you saw, you have two tools to make these functions work legit"}, {"start": 2949.16, "end": 2957.16, "text": " One is to make certain arguments static, and sometimes you'll have to use NumPy functions instead of Jack's functions"}, {"start": 2957.16, "end": 2963.16, "text": " So don't ask me why, I'm only a couple of steps ahead of you, I recently started learning Jack's myself"}, {"start": 2963.16, "end": 2975.16, "text": " But let's try and run this, and this time it works, because we have the mprod will return a concrete value and not the traced object"}, {"start": 2975.16, "end": 2985.16, "text": " And that's it, if you wish, as I said, I'm going to share the colab, I'm going to push it to my github, so go ahead and play with these failure cases yourself"}, {"start": 2985.16, "end": 2997.16, "text": " To better understand how JIT works, but basically just keep in mind that JIT passes in these abstract objects that have the shape information and the data type information and no value"}, {"start": 2997.16, "end": 3001.16, "text": " And that will help you save yourself from a lot of headache, I guess"}, {"start": 3001.16, "end": 3020.16, "text": " Ah, okay, this is something I need to cover because there are some gotchas, some idiosyncrasies of Jack's which you need to understand to be able, so that we'll have easier time in the later videos building up neural networks, etc."}, {"start": 3020.16, "end": 3024.16, "text": " Once you understand a couple of these gotchas, things are going to be way easier"}, {"start": 3024.16, "end": 3037.16, "text": " So, let's start with gotcha number one, pure functions, so Jack's is designed to work only on pure functions, we saw some glimpse of this while we were using print functions a couple of minutes ago"}, {"start": 3037.16, "end": 3049.16, "text": " 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"}, {"start": 3049.16, "end": 3064.16, "text": " 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 the inputs, you just retrieve the outputs"}, {"start": 3064.16, "end": 3071.16, "text": " So that's an informal definition of a pure function and I think it's going to be good enough to understand the following cells"}, {"start": 3071.16, "end": 3087.16, "text": " 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 this return statement here"}, {"start": 3087.16, "end": 3105.16, "text": " So let me call this function, okay, on the first call, again, we have the tracing, so we call the executing function print statement, but the second call, due to the fact that we are passing the same shape and same data type, that's just a float 32 here"}, {"start": 3105.16, "end": 3122.16, "text": " We just use the cache function and we return five, finally, the third call, because we changed the shape and the type here, so we basically have a ray now, it's going to basically call and trigger the tracer again, and I think that's pretty clear at this point of time"}, {"start": 3122.16, "end": 3143.16, "text": " 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 program paradigm, this is probably not a good idea to do because, yeah, unexpected things can happen with your program"}, {"start": 3143.16, "end": 3163.16, "text": " So what's happening here is that, as I said, we are violating both one and two because, 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"}, {"start": 3163.16, "end": 3180.16, "text": " 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 what will happen here, once we call the jitted version of the function the first time, it's going to cache the value of g, which is zero"}, {"start": 3180.16, "end": 3197.16, "text": " Then we are going to update it, and then we are going to call it with five, and it's going to call the cached version of the function, which, if you recall, has g equals to zero somewhere in the Jackspur, and that's why we are going to get the wrong results"}, {"start": 3197.16, "end": 3214.16, "text": " So let me execute this cell, so on the first call we get four, and that's fine, because we passed four, we had g equals zero, that's fine, now g is equal to ten, and we have five here, so we expect fifteen, but we get five because the g version of zero is cached in the function"}, {"start": 3214.16, "end": 3227.16, "text": " And finally, if I call with a different, we change, instead of float, we pass the, like, a Jacks array, this time we trigger the execution and we get the correct result again"}, {"start": 3227.16, "end": 3234.16, "text": " So in any case, this is a wrong idea, a bad idea to add these types of impurities"}, {"start": 3234.16, "end": 3245.16, "text": " Example number three, so Haiku Flux are basically built upon this idea, the whole idea is that it's fine to have a stateful functionality inside of your pure function"}, {"start": 3245.16, "end": 3265.16, "text": " So what I mean by that is the following, so we have x here, and what we do here, 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"}, {"start": 3265.16, "end": 3281.16, "text": " 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 pure, because we're only using x, we're only assigning x to this state"}, {"start": 3281.16, "end": 3287.16, "text": " And we're only outputting whatever came out from this function, so we're not accessing some global variable or whatnot"}, {"start": 3287.16, "end": 3295.16, "text": " So because of that, all of this will be fine, and if I execute this, we're going to get a correct result, and that's 50, okay"}, {"start": 3295.16, "end": 3308.16, "text": " Finally, a fourth example, iterators, since they are stateful, are a no-no, so if you try and use, we built array here, and if we just do like this, if we don't have any iterators, it's going to work"}, {"start": 3308.16, "end": 3322.16, "text": " But on the other hand, if we use the iterator, this is going to fail, even though we semantically would expect the same results, and you can see that's the case, we got 45 here, we got 0 here, even though we should have gotten 45 as well"}, {"start": 3322.16, "end": 3336.16, "text": " Here you can see this lex primitive, so that's one of those mid-level API functions, it's called foriloop, it's basically a smart version of the for loop, which can be later compiled using the XLA"}, {"start": 3336.16, "end": 3347.16, "text": " But the thing I want you to understand here is just that you cannot use iterators because they are stateful, and we are thus violating the purity constraint of Jax"}, {"start": 3347.16, "end": 3356.16, "text": " Okay, that was the gotcha number one, make sure to write pure functions if you want to use them with Jax transforms in general"}, {"start": 3356.16, "end": 3370.16, "text": " Gotcha number two, in-place updates, we saw this already, so you cannot modify the arrays in place, you have to use that at set syntax, so let's see a simple example here"}, {"start": 3370.16, "end": 3378.16, "text": " Once we create the Jax array, we have to use this at set syntax in order to get the output array"}, {"start": 3378.16, "end": 3392.16, "text": " 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"}, {"start": 3392.16, "end": 3411.16, "text": " So I think I mentioned this, so if this seems wasteful, basically don't worry because XLA is smart enough to figure out that if you're not using this input array, then it will not allocate a special memory object for this output array"}, {"start": 3411.16, "end": 3419.16, "text": " It's just going to reuse the input array and modify it in place, even though it does not appear to do so on this higher level perspective"}, {"start": 3419.16, "end": 3438.16, "text": " Don't worry about the expressiveness, we can still do everything we can with NumPy, so here if I create a simple matrix of ones, we can do whatever we want, we can add, we don't just have to assign values to certain locations, we can also do arbitrary operations such as addition"}, {"start": 3438.16, "end": 3451.16, "text": " 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"}, {"start": 3451.16, "end": 3458.16, "text": " So that's cool, that was the second gotcha, let's continue on, the third one is out of bounds indexing"}, {"start": 3458.16, "end": 3476.16, "text": " 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 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"}, {"start": 3476.16, "end": 3492.16, "text": " So when it comes to NumPy, if we allocate an array with ten elements, so from zero to nine, and we try to index into eleventh position, this will throw an exception because there is no eleventh position, there is position zero through nine"}, {"start": 3492.16, "end": 3502.16, "text": " So let me try and run this, and we're going to get exception caught, blah blah blah, basically this is out of bounds for x is zero with size ten"}, {"start": 3502.16, "end": 3521.16, "text": " Okay, let's see what Jax's behavior here is, if we were to try to assign at position eleven, so that's the same example as here, if we try to add twenty-three, it's just going to ignore this operation, so that may be surprising, because it will not throw any error"}, {"start": 3521.16, "end": 3544.16, "text": " So let's see what the result will actually be, so we have no change whatsoever here, finally, if we try to return, to retrieve the element at eleventh position, which does not exist, it's going to clamp eleventh to the last index, which is nine, and so we are going to retrieve nine"}, {"start": 3544.16, "end": 3566.16, "text": " 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 information from you, but as a consequence you get this"}, {"start": 3566.16, "end": 3582.16, "text": " And I guess if you're familiar with NAND behavior, this is kind of similar, because NANDs are also in a sense not acceptable, although we also have non-error behavior, whereas we just get NANDs back instead of the system throwing some exception"}, {"start": 3582.16, "end": 3592.16, "text": " So I mentioned it here, similarly to how invalid floating point arithmetic results in NANDs and not an exception, so keep this in mind"}, {"start": 3592.16, "end": 3608.16, "text": " Gotcha number four, non-array inputs, again this is added by design, it's not a bug, in NumPy if we try to do a sum and we pass a Python list, one, two, three here, we are going to get an expected result, so that's going to be six I guess"}, {"start": 3608.16, "end": 3617.16, "text": " In Jax on the other hand, if we try to do this, we're going to get an exception, so sum requires a NumPy array or scalar arguments got list at position zero"}, {"start": 3617.16, "end": 3631.16, "text": " So why is that? Let's imagine we try and make this 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"}, {"start": 3631.16, "end": 3650.16, "text": " So now let's create a list, X is a simple list, now let's try to make Jaxper to understand what will happen if we called Jit on this function, so let's try and run this, and you'll see that we'll basically have this whole thing unrolled, which is super inefficient"}, {"start": 3650.16, "end": 3665.16, "text": " So 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"}, {"start": 3665.16, "end": 3690.16, "text": " So yeah, keep this in mind, now this is 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, 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 a NumPy or Jax array"}, {"start": 3690.16, "end": 3711.16, "text": " So definitely things like this make me, if I were to recommend to a beginner, I would definitely still recommend PyTorch to be honest over Jax, 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 programs, I think Jax is a really good bet for those guys"}, {"start": 3711.16, "end": 3725.16, "text": " Anyways, let's continue, gotcha number five, random numbers, I mentioned this already, now let's see it in a bit more detail, so basically what we can see here is that NumPy has a stateful pseudo random number generator"}, {"start": 3725.16, "end": 3744.16, "text": " That means that if I execute these two functions in a row, we'll have, so this one will advance the state of the 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 kind of dissect this generator"}, {"start": 3744.16, "end": 3758.16, "text": " So we set the seed to some value, seed is I think set to zero here, and basically we can fetch, there is this function getState, we can fetch the actual state of the NumPy PRNG, and I'm gonna print certain metadata from that state"}, {"start": 3758.16, "end": 3771.16, "text": " Then I'm going to execute, like a sample, a single number, and then we're gonna fetch the state again, and we're gonna sample the number again here, and we're gonna fetch the state again and print it, so let's see what happens here"}, {"start": 3771.16, "end": 3777.16, "text": " So first thing you can notice here is that the numbers here are different because the state is internally advanced"}, {"start": 3777.16, "end": 3791.16, "text": " So what you can notice here is that first of all NumPy is using this thing called Mersenne Twister, hopefully I'm pronouncing it correctly, PRNG, and it's known to have a number of problems, there is a link in their documentation for why that is, it's not that important for us"}, {"start": 3791.16, "end": 3809.16, "text": " What is probably useful to know is that it basically has the state consists of 624 unsigned integers, and every time you sample from that generator, we are basically consuming the entropy of the generator"}, {"start": 3809.16, "end": 3823.16, "text": " Bottom line, what I want you to take out from this is that NumPy's PRNG has some problems, and it's stateful, which is problematic because remember we need to have pure functions, we are in the functional programming world"}, {"start": 3823.16, "end": 3834.16, "text": " 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"}, {"start": 3834.16, "end": 3840.16, "text": " So basically we are now manipulating the state of the PRNG externally as opposed to internally"}, {"start": 3840.16, "end": 3848.16, "text": " So key is simply like a tuple of two unsigned integers, 32-bit integers, so that's the state"}, {"start": 3848.16, "end": 3859.16, "text": " Now what's the trick here? So now if we were to sample using that same state, we are going to get obviously the same results, so that's different compared to the NumPy behavior"}, {"start": 3859.16, "end": 3866.16, "text": " And the reason is again because we are not modifying the state, we are not advancing the state, and that's why we always get the same results"}, {"start": 3866.16, "end": 3873.16, "text": " So again, important to notice here, state is preserved, state has not changed, and secondly, the results are the same"}, {"start": 3873.16, "end": 3881.16, "text": " So what do we do? Obviously this is not a random number, this is a constant function, so how do we get a randomness in Jax?"}, {"start": 3881.16, "end": 3890.16, "text": " So here is a trick, every time you want to create a new random value, you basically just call this split function, and it's going to return key and subkey"}, {"start": 3890.16, "end": 3898.16, "text": " Subkey is used to generate a novel number, and key can subsequently be used again in the split function to get novel key and novel subkeys"}, {"start": 3898.16, "end": 3901.16, "text": " So every time you want to generate a new number, you have to call this splitting"}, {"start": 3901.16, "end": 3914.16, "text": " So this may seem like very rough and problematic, but believe me, it helps solve various issues that are caused by randomness in libraries such as NumPy, etc"}, {"start": 3914.16, "end": 3923.16, "text": " When you try to use them outside of the context for which they were designed for, and that's like basically I guess CPU, single threaded programs, etc"}, {"start": 3923.16, "end": 3931.16, "text": " Okay, so let's run this cell now and see the results. So we have the old key, old key got converted into the new key, so that's the new state"}, {"start": 3931.16, "end": 3937.16, "text": " And we have this subkey which we use to generate the random number"}, {"start": 3937.16, "end": 3944.16, "text": " So a couple of notes here, first things first is that basically you can split into more subkeys than just the two of these"}, {"start": 3944.16, "end": 3954.16, "text": " And secondly, there is no semantic difference between the key and the subkey, these are basically recommendations for how to organize these states"}, {"start": 3954.16, "end": 3965.16, "text": " And basically you use key to generate novels 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"}, {"start": 3965.16, "end": 3977.16, "text": " And that's pretty much it. So after having explained all these complex details about how to handle PRNGs compared to NumPy, is there any good reason for it?"}, {"start": 3977.16, "end": 3988.16, "text": " And the answer is yes. So why this design? And the answer is can the code with the current design, with NumPy's design be reproducible, parallelizable, and vectorizable?"}, {"start": 3988.16, "end": 3998.16, "text": " So in the case of NumPy, number one is obeyed, in the case where you have a single threaded program on the CPU, basically there is no problem"}, {"start": 3998.16, "end": 4009.16, "text": " So let's see this concrete example. Let's assume we have a function called bar that generates a random number and a function called baz that also generates a random number"}, {"start": 4009.16, "end": 4016.16, "text": " Finally we have a function foe, and I don't know if I'm pronouncing these right, so this function returns bar plus two times baz"}, {"start": 4016.16, "end": 4026.16, "text": " And I don't know if you can see the problem with this code, once we start and try to call this foe function on line 14"}, {"start": 4026.16, "end": 4039.16, "text": " Now the problem is NumPy assumes a single threaded environment, and basically Python guarantees that this will be executed, to the best of my knowledge, from left to right"}, {"start": 4039.16, "end": 4049.16, "text": " Which means every time we run this, we are going to get the same result back, which means number one is obeyed, so that means we have reproducible programs"}, {"start": 4049.16, "end": 4065.16, "text": " So now what happens if we get this function, or if we, and the JIT decides to basically parallelize this program and calls bar on one accelerator, on one core, and baz on another one"}, {"start": 4065.16, "end": 4072.16, "text": " So what can happen is that the order of execution can change, and if that happens we may get different results"}, {"start": 4072.16, "end": 4088.16, "text": " So if the first function to be called returns 0.3 and the second one returns 0.4, depending on the order we'll either have 0.3 plus 2 times 0.4, or we'll have basically 0.4 plus 2 times 0.3"}, {"start": 4088.16, "end": 4100.16, "text": " And that basically leads to different outcomes, which means this will not be a reproducible result in the case of parallel, like computing on parallel, like cores, machines, whatever"}, {"start": 4100.16, "end": 4113.16, "text": " Similarly for the vectorization part, basically NumPy guarantees that if you generate numbers, if you do this, if you iterate in a for loop and you generate sequentially random numbers"}, {"start": 4113.16, "end": 4122.16, "text": " That will give you the same results as if you were to just set size to 3, on the other hand, Jax does it differently, if you generate individually"}, {"start": 4122.16, "end": 4129.16, "text": " That will give you different results compared to if you generate from all the numbers at once using a specific key"}, {"start": 4129.16, "end": 4138.16, "text": " So, and here again you can see an example where we generate three subkeys using the split function and not only two keys, as I previously mentioned that"}, {"start": 4138.16, "end": 4149.16, "text": " So let's run this function and you can see that basically NumPy has the same results, whereas Jax does not have the same results"}, {"start": 4149.16, "end": 4155.16, "text": " So if you're using the common SIMD pattern, so that's the single instruction multiple data, which is common in machine learning"}, {"start": 4155.16, "end": 4165.16, "text": " Whereas you want to apply the same functionality across different batches, you sometimes want to have the same randomness being applied across all of the batches"}, {"start": 4165.16, "end": 4175.16, "text": " And NumPy does not allow to do that, that's how I understood this part about the violating the vectorization problem, let me know if I got this wrong"}, {"start": 4175.16, "end": 4183.16, "text": " Okay, that's pretty much it, that was the random numbers, we have just two more gotchas to go and that's it"}, {"start": 4183.16, "end": 4193.16, "text": " So gotcha number four is the control flow, it's fairly, we've seen something similar to this, so basically control flow plus grad, nothing, there is no problem"}, {"start": 4193.16, "end": 4200.16, "text": " Basically this transform function, so the grad transform function can deal with these types of conditioning on the value of a function"}, {"start": 4200.16, "end": 4213.16, "text": " I went ahead and ran this and you can see the results, basically, yeah, just analyze this function, you'll see that at three there is this jump where we have a piecewise defined function here"}, {"start": 4213.16, "end": 4225.16, "text": " The whole point is if we were to take a grad of this function and evaluate it at two and at four, so that's just before this discontinuity and just after it, we'll get valid gradients"}, {"start": 4225.16, "end": 4236.16, "text": " 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 are conditioning on x"}, {"start": 4236.16, "end": 4244.16, "text": " So we have to, the solution is to make it a static argument and then it's gonna work, so if I run this, it's going to work"}, {"start": 4244.16, "end": 4250.16, "text": " And we already saw how to handle this case of conditioning on a value, so let's see some more interesting cases"}, {"start": 4250.16, "end": 4259.16, "text": " So here we are conditioning again on the value, but this time the value decides the length, the number of loops we'll have in this function"}, {"start": 4259.16, "end": 4275.16, "text": " And the way to go around this is again make this a static argument which will make jit trace this function using the forex, using the abstract tracer, so the shape and data type"}, {"start": 4275.16, "end": 4280.16, "text": " Whereas for n is gonna use a concrete value, so let me run this cell and see what we get"}, {"start": 4280.16, "end": 4292.16, "text": " Okay, so as you can see we have some huge, huge checks per here and the reason being is because we have, as you can see here, 15 was passed for n"}, {"start": 4292.16, "end": 4300.16, "text": " And this is the best way that jit can deal with these types of primitive, of native Python for loops"}, {"start": 4300.16, "end": 4310.16, "text": " Again, importantly, you should not change static values too often, otherwise we'll be triggering recompilation all of the time"}, {"start": 4310.16, "end": 4315.16, "text": " And then the overhead will maybe be detrimental to the speed of your application"}, {"start": 4315.16, "end": 4319.16, "text": " Let's see how this can be avoided and make a bit more optimal"}, {"start": 4319.16, "end": 4328.16, "text": " So a better way to do this is to use the low level API again, so the lex API, and there is this, as we saw, for iLoop function"}, {"start": 4328.16, "end": 4335.16, "text": " And we can rewrite the above problem and I went ahead and just rewrote it here and you can see that"}, {"start": 4335.16, "end": 4343.16, "text": " So this is the same function as the above one just using the lex API and if we now do the make, if I call the make checks per"}, {"start": 4343.16, "end": 4352.16, "text": " And 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"}, {"start": 4352.16, "end": 4357.16, "text": " I haven't profiled this example but yeah, I can assume it's more efficient"}, {"start": 4357.16, "end": 4367.16, "text": " You can go ahead and analyze 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"}, {"start": 4367.16, "end": 4371.16, "text": " You can sometimes get more out of it"}, {"start": 4371.16, "end": 4379.16, "text": " Finally, the only reason I have this one is to understand that you can condition sometimes on the dimensionality of your data"}, {"start": 4379.16, "end": 4387.16, "text": " And that's allowed because imagine, this means that for the whole class of x where you have, where x is two dimensional"}, {"start": 4387.16, "end": 4393.16, "text": " This can, like a function can be cached and will work and we can fetch it and use it"}, {"start": 4393.16, "end": 4396.16, "text": " So let's run this one and let's see the results"}, {"start": 4396.16, "end": 4403.16, "text": " So I passed like an input array which is not two dimensional and because of that we took this branch"}, {"start": 4403.16, "end": 4407.16, "text": " Which means whatever we get as an input we just return as the output"}, {"start": 4407.16, "end": 4414.16, "text": " And that's why we have this super simple like a Jacksper for this particular function being traced with this input"}, {"start": 4414.16, "end": 4422.16, "text": " Hopefully all of these helped you crystallize and understand JIT because I guess, I think that's arguably the hardest part to understand about JAKS"}, {"start": 4422.16, "end": 4424.16, "text": " How JIT works in the background"}, {"start": 4424.16, "end": 4432.16, "text": " That's it, final, final, final gotcha is like how to handle NANDs in JAKS"}, {"start": 4432.16, "end": 4439.16, "text": " The usual non-error behavior is to simply return a NAND in the case of like operations such as this one, division by zero"}, {"start": 4439.16, "end": 4443.16, "text": " So if I were to comment out this thing here and let me run this cell"}, {"start": 4443.16, "end": 4447.16, "text": " So we won't have any error, we'll just have a NAND"}, {"start": 4447.16, "end": 4452.16, "text": " So if we want to actually debug and understand where the NANDs came from"}, {"start": 4452.16, "end": 4458.16, "text": " You wanna like the program to throw some exceptions, you can do, for example, this, you can find more in the docs"}, {"start": 4458.16, "end": 4463.16, "text": " But basically there is a way to do it and now it will be throwing exceptions, I guess"}, {"start": 4463.16, "end": 4465.16, "text": " Let me try again"}, {"start": 4465.16, "end": 4468.16, "text": " Yeah, now it's throwing an exception"}, {"start": 4468.16, "end": 4470.16, "text": " So yeah, let's keep that in mind"}, {"start": 4470.16, "end": 4474.16, "text": " Final cell, JAKS enforces single precision"}, {"start": 4474.16, "end": 4484.16, "text": " 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"}, {"start": 4484.16, "end": 4486.16, "text": " So it means 8 bits only"}, {"start": 4486.16, "end": 4493.16, "text": " And because NAMPA is aggressive in promoting like certain variables into double, so that means 64 bits"}, {"start": 4493.16, "end": 4497.16, "text": " JAKS made it by design that they are enforcing 32 bits"}, {"start": 4497.16, "end": 4502.16, "text": " So that may lead to some, again, some peculiar behavior such as this one"}, {"start": 4502.16, "end": 4511.16, "text": " You say I wanna like a vector with thousand random numbers which are 64 like bit long"}, {"start": 4511.16, "end": 4517.16, "text": " And if we were to print the data top of the actual array, we'll have flow 32, which is not intuitive, right?"}, {"start": 4517.16, "end": 4524.16, "text": " 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"}, {"start": 4524.16, "end": 4532.16, "text": " As a quick summary, we saw the basics of JAKS such as the fact that it's using, it's based on functional programming paradigm"}, {"start": 4532.16, "end": 4538.16, "text": " We saw like various transform functions such as JIT, GRAD, DMAP"}, {"start": 4538.16, "end": 4546.16, "text": " Then we saw, we went deeper into JAKS, we saw the layered, like, layer API"}, {"start": 4546.16, "end": 4554.16, "text": " We saw many details of how JIT works and we saw many gotchas that JAKS has which may catch you by surprise"}, {"start": 4554.16, "end": 4557.16, "text": " So you should be aware and cognizant of these"}, {"start": 4557.16, "end": 4564.16, "text": " Finally, some conclusion of mine to who should be using JAKS and who should be using other frameworks"}, {"start": 4564.16, "end": 4571.16, "text": " As I said, I think JAKS is very good if you're a researcher, you wanna be very flexible and you wanna have a powerful tool"}, {"start": 4571.16, "end": 4578.16, "text": " And you also wanna, once you train big models such as in whatever industry I lab you're working, you're gonna be training big models"}, {"start": 4578.16, "end": 4582.16, "text": " So having like very optimized code is super important"}, {"start": 4582.16, "end": 4590.16, "text": " 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 JAKS"}, {"start": 4590.16, "end": 4592.16, "text": " And I don't think that's going to change"}, {"start": 4592.16, "end": 4597.16, "text": " So basically all of these optimization details such as the fact you cannot use Python list"}, {"start": 4597.16, "end": 4601.16, "text": " That will just throw exception because it can hurt the performance"}, {"start": 4601.16, "end": 4606.16, "text": " Because of all of that and because of the functional paradigm"}, {"start": 4606.16, "end": 4610.16, "text": " I think it's easier to just start with PyTorch"}, {"start": 4610.16, "end": 4612.16, "text": " Especially when we get to neural networks"}, {"start": 4612.16, "end": 4620.16, "text": " It's a more natural approach to approach neural networks with object oriented, from the object oriented perspective"}, {"start": 4620.16, "end": 4626.16, "text": " At least in my opinion maybe I lack experience with JAKS so yeah, but that seems to be the case"}, {"start": 4626.16, "end": 4632.16, "text": " Anyways, in the next video we're going to cover some concepts such as PyTrees, handling states"}, {"start": 4632.16, "end": 4640.16, "text": " Which are all basic components that we'll need to later build neural networks from scratch in PureJAKS and also in Haiku or Flex"}, {"start": 4640.16, "end": 4646.16, "text": " Hopefully you found this video useful, it took a lot of time to prepare the notebook, to prepare everything I wanted to show you here"}, {"start": 4646.16, "end": 4651.16, "text": " So if you did like it, consider sharing it out, also consider subscribing to this channel if you haven't already"}, {"start": 4651.16, "end": 4659.16, "text": " Also do join the Discord community, there is a lot of smart people there and we have a community for, like there is more than 1,000-100 people at this point of time"}, {"start": 4659.16, "end": 4663.16, "text": " So somebody will help you out if you have some doubt or questions"}, {"start": 4663.16, "end": 4676.16, "text": " In any case, until next time, bye bye!"}]