JAX in 100 Seconds

Key Takeaways

Jax it's just another accelerated linear algebra library but capable of mega fast numerical Computing on futuristic new hardware to understand Jax let's start with the x or accelerated linear algebra in Python there's already a great linear algebra Library called numpy and Jax is nearly identical to numpy it allows you to create multi-dimensional arrays then do scientific Computing with them like add them together or get the dot product however Jax enforces some constraints that numpy does not have like immutable arrays and pure functions which allows it to automatically compile to low-level code that can run on accelerated Hardware like gpus and tpus the a stands for autograd Jax was developed by Google along with team members from the original autograd Library you see virtually every facet of machine learning requires some calculus you'll need to compute gradients for optimization algorithms and back propagation and neural networks and autograd allows you to automatically differentiate python functions or in simple terms it calculate the rate of change of a function based on its inputs then finally the J stands for just in time compilation when you write a function in Jax it's transformed into a primitive set of operations these so-called Jack spers are lazily compiled and can be evaluated like a mini functional programming language to get started install Jacks either for your CPU GPU or TPU and now we can start doing high performance array Computing let's first create a couple of two-dimensional arrays now in numpy if I wanted to change the values in an array I could mutate the data structure directly but in Jax that's an error because arrays are immutable instead we use the at function to select an index then set to change the value and return a new array from there we can do some linear algebra like element wise addition element wise multiplication or get the dot product of the two arrays what's really cool about Jack though is automatic differentiation imagine you're building a high yield nuclear warhead in your mom's basement you'll likely write a python function that looks like this that calculates the height of the mushroom cloud based on the amount of time after detonation this will give us the height at any point in time but we also want to know the instantaneous rate of change or how fast the cloud is growing at a given point in time with Jax we can differentiate this function by passing our function to Jack's grad which returns a new function that computes the derivative but if we modify this function to take an array of parameters the end result will be a gradient which is an array of partial derivatives giving us the rate of change with respect to each input variable on the original function because grad returns a function we can even apply it to its own output to differentiate again and get higher order derivatives but more importantly you can use math just like this in Jax to build the future of machine learning gradients tell us how to adjust model parameters to decrease the loss or improve the performance of a model and you can start building deep neural networks with it right now using libraries like flax this has been Jax in 100 seconds but if you truly want to Master machine learning you'll also need to understand computer science calculus and statistics you can start making that happen today for free thanks to this video sponsor brilliant not only does their platform have a massive collection of content related to these topics but their fun fun Hands-On exercises will help you actually retain what you study just like this video every lesson is designed to be concise and rewarding making complex topics in math and science approachable for everyone you can try out everything brilliant has to offer for free for 30 days by visiting brilliant.org fireship or scan this QR code for 20% off their premium annual subscription thanks for watching and I will see you in the next one