Implementación de núcleos computacionales para redes neuronales con JAX

Abstract

[ES] Esta investigación se realiza con un objetivo principal que es el estudio de la portabilidad de la biblioteca JAX de Python, desarrollada por Google, en diferentes componentes hardware (CPU, GPU y TPU). JAX es una herramienta de gran alcance que combina la familiaridad y simplicidad de NumPy con avanzadas capacidades de diferenciación automática y ejecucián de alto rendimiento hardware acelerado. Además del estudio de portabilidad de código llevamos a cabo una comparativa de rendimiento con otras bibliotecas populares como son NumPy y Cupy. Esta última solo para GPU, utilizando la multiplicación de matrices como caso de estudio. La comparativa se ha realizado en términos de GFLOPS. Las pruebas de rendimiento se llevaran a cabo en Notebooks de Google Colab y Kaggle, aprovechando de esta forma recursos hardware diversificados. Se obtiene y analiza el rendimiento máximo alcanzable, en segundos, para cada componente hardware.

[EN] This research is conducted with a primary objective: to study the portability of the JAX library for Python, developed by Google, across different hardware components (CPU, GPU, and TPU ).JAX is a powerful tool that combines the familiarity and simplicity of NumPy with advanced capabilities for automatic differentiation and high-performance hardware- accelerated execution. In addition to the code portability study, we conduct a performance comparison with other popular libraries such as NumPy, Numba and CuPy, using matrix multiplication as the case study. The comparison is performed using billions of floating point operations per second (GFLOPS ). The performance tests will be carried out in Google Colab and Kaggle notebooks, thereby leveraging the diversified hardware resources. The maximum achievable performance, in seconds, for each hardware component is obtained and analyzed.