JIT Compiling CPython with Numba & JAX
Tuominen, Joonas (2025)
2025
Tieto- ja sähkötekniikan kandidaattiohjelma - Bachelor's Programme in Computing and Electrical Engineering
Informaatioteknologian ja viestinnän tiedekunta - Faculty of Information Technology and Communication Sciences
Hyväksymispäivämäärä
2025-11-06
Julkaisun pysyvä osoite on
https://urn.fi/URN:NBN:fi:tuni-2025110610444
https://urn.fi/URN:NBN:fi:tuni-2025110610444
Tiivistelmä
Python, or more precisely CPython, is one of the most common programming languages. It is widely used in varied applications. However, with the recent focus on machine learning together with the advent of CPython machine learning libraries, such as Tensorflow, Keras and PyTorch, there has been need to optimize CPython execution to cope with these computationally heavy workloads. This has in turn sparked the development of just-in-time (JIT) compilers for use with these machine learning libraries. Even though these JIT compilation libraries have been mainly designed for machine learning applications, their use is not strictly bound to only this domain. Thus, they could be potentially utilized in optimizing the CPython runtime via JIT compilation.
This thesis aims to showcase and compare two CPython JIT compiler libraries: Numba and JAX. This is accomplished by first discovering the current performance shortfalls of the CPython interpreter, exploring the innerworkings of these compilers, as well as comparing their limitations. Finally, these libraries are compared by their relative execution times and ease of use.
According to the findings, the more machine learning centric JAX clearly shines in situations, where the used functions are more dependent on linear algebra and more closely follow the functional programming paradigm. Also, JAX seems to be able to better optimise the executed computations. JAX however requires an extremely specific modus operandi and thus may often lead to struggles with integrating the JAX JIT compiler with the existing code base. Numba on the other hand is rather relaxed and thus, is extremely easy to use. In addition to this, Numba is often able to significantly improve the execution runtime. However, using Numba may cause unexpected behaviour, as in testing its usage seems to have caused a performance bottleneck in one of the performance tests. It should however be noted, that, the used testing was limited and further testing is necessary to validate the findings.
This thesis aims to showcase and compare two CPython JIT compiler libraries: Numba and JAX. This is accomplished by first discovering the current performance shortfalls of the CPython interpreter, exploring the innerworkings of these compilers, as well as comparing their limitations. Finally, these libraries are compared by their relative execution times and ease of use.
According to the findings, the more machine learning centric JAX clearly shines in situations, where the used functions are more dependent on linear algebra and more closely follow the functional programming paradigm. Also, JAX seems to be able to better optimise the executed computations. JAX however requires an extremely specific modus operandi and thus may often lead to struggles with integrating the JAX JIT compiler with the existing code base. Numba on the other hand is rather relaxed and thus, is extremely easy to use. In addition to this, Numba is often able to significantly improve the execution runtime. However, using Numba may cause unexpected behaviour, as in testing its usage seems to have caused a performance bottleneck in one of the performance tests. It should however be noted, that, the used testing was limited and further testing is necessary to validate the findings.
Kokoelmat
- Kandidaatintutkielmat [10747]