Software

When we communicate scientific results, we use a standardised format. Can we do the same for the implementations of such analyses?

joint

Code for building and fitting diverse multi-modal generative model architectures. You may ask: what is the point of this now? This is a fair question. That said, the machine learning literature develops the components of multi-modal models (e.g. attention and generative models for PDFs) in vastly different directions.

This software allows you to swap modules in such architectures around, with various encoder/decoder modules for different modalities. This is useful e.g. in the case I had where I swapped in a flow map architecture in place of a diffusion model, with fast-attention modules for the DiT module that parameterised the flow map velocity model.

See here.

sbgm

A code for fitting score-based diffusion models in jax.

See here.

pip install sbgm

sbiax

This is a code for implementing your own SBI experiments and analyses.

See here.

stak

This is a code for running large hierarchical models (chains of generative models or a forward model) on sharded GPU configurations (I need to extend this to multi-host processes!).

It could be useful with JAX’s ffi interface, where foreign C++ codebases could be ported for use on multi-GPU clusters.

See here.