Evaluation of JAX in AI/ML software engineering

From ISLAB/CAISR
Title Evaluation of JAX in AI/ML software engineering
Summary Analysis of the benefits of JAX (and/or similar solutions) in terms of performance, development time, module reusability, etc.
Keywords
TimeFrame Fall 2023
References
Prerequisites
Author
Supervisor Veronica Gaspes, Slawomir Nowaczyk
Level Master
Status Open


Many new tools are being created with the goal of making it easier to build AI/ML solutions. One such tool is JAX:
https://jax.readthedocs.io/en/latest/

This thesis is about the evaluation of JAX, taking multiple different perspectives into account. One aspect is of course performance, but there are also "bigger" software engineering questions: can it be used to express things in a more familiar (mathematical) way than e.g. TensorFlow? And thus allow for shorter development time, more reusable modules, the possibility of exploring different methods for the same problem, etc.?

Please note: this topic is just an early idea, not fully developed. If you are interested in it, contact Veronica & Slawomir before the deadline to discuss how to make it more concrete.