JAX (Python)
What is JAX (Python)?
JAX (Python) stands for "Just After eXecution" and is a newly developed library by DeepMind that focuses on deep learning as well as machine learning. JAX, in contrast to the well-known Tensorflow is not an official product of Google and is becoming increasingly popular due to its NumPy-like syntax, especially in the research sector.
JAX uses the maximum number of FLOPs to generate optimal code. Viewed precisely, this library is actually a just-in-time compiler, or JIT compiler for short. The just-in-time principle means that JAX can be used not only on the CPU, but also on the GPU or TPU, which is another major advantage and at the same time the main feature of JAX. JAX can automatically distinguish between native Python and NumPy code, as well as differentiate subsets and take multiple derivations.
In which area is JAX used?
The use of such software can be found Increased research and development again. Mainly due to the complex working structure of this software, other and partly new solutions can be calculated. However, the changeover to this new principle is not entirely easy, which also becomes clear later in the comparison with PyTorch.
In principle, one can speak here of a similar field of application as with PyTorch, but due to its complex form, not every company uses JAX. At least not at the moment.
Unsurprisingly, the largest deployment comes from the self-employed IT technician sector. Especially for programming and structuring or the calculation of graphics cards, JAX can make the difference. It is easy to explain why the volume is so high in this sector in particular. New companies start directly with this software and not with the predecessors, which eliminates the need for later retraining or adaptation. In this case, the software that is not officially supported by Google is used directly. Looking to the future, it can be deduced that the use and thus the spread of this software will continue to increase in the future.
What is the difference between JAX and PyTorch?
The main difference between the two programmes is the basic structure of calculations that are to be used for a project. This results in a completely different structure, which thus also has an impact on the tracing of the calculated data. This can be explained as follows:
PyTorch builds a graph during the run. Here, a call to backward() can be made on a result node. This results in an extension to each intermediate node in the graph by the gradient of the result node, which is related to this intermediate node. This offers the possibility to directly influence various individual elements that are created / calculated during the run.
With Jax, on the other hand, the calculation is expressed as a Python function, which becomes a gradient function by converting grad(), which can then be evaluated in the same way as the calculation function. In addition, the output is not provided as a typical output, but as a gradient of the output, which refers to the first parameter of the function.
This difference in particular makes it clear that completely different codes and models have to be written in both variants. This is why one speaks of a major change when switching from one programme to the other.
The greatest commonality of the two types of software is their flexibility and the associated wide range of applications where they can be found. In addition, both work in the low and high-level API. Both types are researchers, which can of course also be seen as a commonality.
