Kolen Cheung

Research Engineer | Physics & Bayesian Inference | AI from first principles

Kolen Cheung trained as a physicist — PhD from UC Berkeley, CMB cosmology, Bayesian inference at petabyte scale — but has always been more at home in mathematics. The instinct to derive before trusting, to find the structure underneath, runs through everything he does.

His path to JAX came through scientific computing rather than machine learning. What drew him was a specific insight: JAX enforces pure functional programming while remaining performant for numerical work, a combination that had not previously existed. Code can be structurally isomorphic to the mathematics it implements — defining what, not how.

Alongside physics, he studies biblical languages: grammar, semantic domains, how meaning transfers across translation. That lens — asking what a system is actually doing with language — is part of what makes the current generation of AI genuinely interesting to him.

He is currently working through the mathematics of neural networks from first principles.