JAX As Accelerated NumPy — JAX documentation

jax.readthedocs.io · 810 words · saved by 1 readers

Fundamentally, JAX is a library that enables transformations of array-manipulating programs written with a NumPy-like API. Over the course of this series of guides, we will unpack exactly what that means. For now, you can think of JAX as differentiable NumPy that runs on accelerators. The code below shows how to import JAX and create a vector. So far, everything is just like NumPy. A big appeal of JAX is that you don’t need to learn a new API. Many common NumPy programs would run just as well in JAX if you substitute np for jnp. However, there are some important differences which we touch on at the end of this section. You can notice the first difference if you check the type of x. It is a variable of type DeviceArray, which is the way JAX represents arrays. One useful feature of JAX is that the same code can be run on different backends – CPU, GPU and TPU. We will now perform a dot product to demonstrate that it can be done in different devices without changing the code. We use %timei

; } .ethical-fixedfooter .ethical-callout { color: #999; padding-left: 6px; white-space: nowrap; } .ethical-fixedfooter a, .ethical-fixedfooter a:hover, .ethical-fixedfooter a:active, .ethical-fixedfooter a:visited { color: #404040; text-decoration: none; } .ethical-fixedfooter .ethical-close { position: absolute; top: 0; right: 5px; font-size: 20px; line-height: 20px; } /* RTD Theme specific customizations */ .wy-nav-side .ethical-rtd { /* RTD theme doesn't correctly set the sidebar width */ max-width: 300px; padding: 0 1em; } .ethical-rtd .ethical-sidebar { /* RTD theme doesn't set sidebar t

JAX As Accelerated NumPy — JAX documentation

related reading