Context
Starting in 2029, the Large Hadron Collider (LHC) at CERN will enter its High-Luminosity phase, increasing collision rates by more than an order of magnitude. This unprecedented data flow will require new, highly efficient particle reconstruction algorithms capable of combining precision and speed in complex detector environments such as ATLAS.
Graph Neural Networks (GNNs) have demonstrated strong potential for this task. However, deploying them at scale demands optimized inference performance. The L2IT laboratory plays a key role in developing a GNN-based reconstruction prototype for the ATLAS experiment, in collaboration with international research partners.
Objectives
The goal of this internship is to reimplement an existing PyTorch-based GNN model in JAX and benchmark its inference performance across various hardware architectures.
Key tasks will include:
- Reproducing or transferring pre-trained weights to JAX
- Validating numerical consistency and reconstruction quality
- Evaluating inference speed, scalability, and efficiency
- Comparing JAX performance with a TensorRT-optimized PyTorch baseline
This study will provide valuable insights into the suitability of JAX for high-performance GNN inference in large-scale physics applications.
Skills & Profile
Required:
- Strong proficiency in Python
- Solid experience with PyTorch or other deep learning frameworks
Appreciated:
- Familiarity with JAX, machine learning acceleration, or graph-based models
Learning Outcomes
The student will gain hands-on experience with JAX, deepen its understanding of advanced neural network architectures such as GNNs, and work in close collaboration with international partners, including researchers from the Lawrence Berkeley National Laboratory.
Practical Information
Location: L2IT Laboratory, Toulouse, France
Duration: 6 months
Start date: Flexible, between January and March 2026
Eligible candidates: Master’s students in Computer Science, Applied Mathematics, or Engineering, with a strong background in Python and deep learning frameworks
