Previous talks at the SCCS Colloquium

Fritz Hofmeier: Integrating Numerical Backend Modularity into Torchquad Using Autoray

SCCS Colloquium |


GPU programming for the Python3 language has become increasingly popular, especially in the context of machine learning. However, the number of backends which provide an array programming interface for GPUs has increased as well and algorithms which are implemented with different interfaces are often incompatible. torchquad is an example of a Python3 module which implements multi-dimensional numerical integration on GPUs and it has employed the array programming interface provided by the backend PyTorch. No module providing equivalent features is available for other backends, such as JAX and TensorFlow.
The autoray module offers an interface for array programming that automatically works for numerical Python3 modules with an API which is similar to NumPy's API. Since it enables algorithm implementations which support multiple backends at once, autoray tackles the compatibility problem.

The topic of this thesis is the adaptation of torchquad's code so that it uses autoray instead of PyTorch and thus supports numerical integration with the backends NumPy, JAX and TensorFlow in addition to PyTorch. Furthermore, this thesis includes investigations about the limitations of modularity between backends, code compilation with PyTorch, JAX and TensorFlow, and performance differences between backends.
Execution time comparisons reveal that PyTorch, JAX and TensorFlow all parallelize well on the GPU, and without compilation, PyTorch is often more than twice as fast as JAX and TensorFlow. Furthermore, the use of autoray has not led to a decrease in code readability in torchquad and the limitations of modularity are mostly due to backend-specific differences, so we can conclude that autoray is in general well-suited to implement computations with support for multiple backends.

Master's thesis presentation (Informatics). Fritz is advised by Dr. Pablo Gómez and Fabio Gratl.