Doing Post-Quantum Cryptography in JAX

In 2018 while studying ML and cryptography around the same time, I realized that many cryptographic algorithms can be expressed as computation graphs, the same ones supported by major ML frameworks, which led to a completely frivolous attempt to implement cryptographic algorithms in TensorFlow, just to see if it would work.

The world has changed a lot since 2018. JAX is growing in popularity as an ML framework, and in the cryptography space, post-quantum cryptography has gone from a mostly theoretical threat to a slightly more real one. So continuing the motivation of the original post, I decided to look into implementing a post-quantum cryptography algorithm in JAX.