Many protein design pipelines implemented in Python and based around PyTorch or JAX, limiting users to use models defined in the corresponding framework while being purpose-build for a specific computational infrastructure (CUDA, essentially). On the other hand, the Julia programming language aims to directly address this two-language problem: it allows for flexible prototyping as in Python while also enabling users to write fast, optimised code at C++ level speed. Moving code to specific devices is as simple as running `to_device(x)` and multiple dispatch takes care of running device optimised code. This means that all code can be written in pure Julia without requiring additional frameworks to improve execution speed. Recently, Reactant.jl has also enabled users to directly compile Julia functions to XLA-optimised code, meaning generic Julia code can directly compete with JAX without needing any code changes.

This project focusses on implementing a protein design pipeline in Julia that supports all GPU infrastructures (including TPU through Reactant.jl) and can call any ML model through the Open Neural Network Exchange (ONNX) framework (including Python models). PyTorch or JAX models can also be compiled down to StableHLO to support native Julia gradient estimation using Enzyme.jl. This means that the ecosystem does not require users to port their Python models to Julia before they can start using the pipeline.

Furthermore, the project aims to implement a high-level programming language for molecular design, allowing users to select any number of binding targets, provide specific design constraints like length, solubility or symmetry, as well as create highly tailored Energy functions from a set of predefined or custom loss terms to conditionally optimise their design for specific features. The language will be design-agnostic, meaning that users can re-use existing pipelines no matter what they are optimising (e.g. amino acid sequences or molecular graphs).

In parallel, I have been working on porting popular structure prediction models (such as Boltz2, AlphaFold and Protenix) to Julia. The benefits here are two-fold: multiple dispatch ensures that optimised code is run for each specific platform, while the LLVM compiler generates optimised code that is capable of improving run times over equivalent Python code. For example, for specific compute-heavy layers in AlphaFold2 (including Triangle Attention layers) we found that Julia implementations resulted in 1.1-1.5x improvements in run time on CPU and Metal GPU. Similarly, early results in LuxProteinMPNN.jl depicts nearly a 2x speed-up over the original ProteinMPNN PyTorch model on Metal hardware when sampling amino acid sequences based on 3D structure coordinates. More extensive benchmarking is currently planned.