Code
Selected research software for scientific computing and 21-cm cosmology.
Two of these active research repositories are currently private. Public releases will follow when the associated work is ready.
JAXMg
Open source · Python · JAX · C++/CUDA · XLA FFI
JAXMg brings distributed dense linear algebra to JAX by connecting JIT-compiled programs to NVIDIA's cuSOLVERMp library. It provides multi-GPU and multi-node linear solves and matrix decompositions for problems exceeding a single GPU's memory. I co-developed the package and its distributed execution model.
GitHub · Documentation · Paper
REACH GPU Data Analysis Pipeline
Private · active research · Python · JAX · BlackJAX · HEALPix · GPU/HPC
An end-to-end simulation and Bayesian inference pipeline for global 21-cm cosmology. I developed the JAX and BlackJAX implementation to combine sky, foreground, signal, beam, and noise models with GPU-accelerated nested sampling, enabling large validation studies and evidence-based model comparison.
BayesEoR-JAX
Private · active research · Python · JAX · JAXMg · Bayesian inference · radio interferometry
A JAX implementation and extension of BayesEoR for estimating the redshifted 21-cm Epoch of Reionization power spectrum from interferometric visibilities. My work adds differentiable, hardware-accelerated forward modelling, JAXMg-backed marginalised likelihoods, and multi-device execution.