JaxProof: JAX code generation and verification

Stream: computer science

[Lun Song]

This is an ongoing project to generate and verify Jax code. Now you can write neural networks in Lean
Screenshot from 2025-12-28 18-50-47.png

[Alex Meiburg]

Very cool! How does this handle the opaqueness of Float operations?

[Lun Song]

Currently I treat float numbers as real numbers, but it can be extended to other approximations like treating float numbers as random variables. The defined jax functions are actually function templates which rely on a typeclass Jax.Impl, which I have defined two instances ,Jax.Tracer to generate code, and Jax.Array to prove properties assuming float is real. To support other approximations or backends, we can simply define a new Jax.Impl instance