Jax-Lean Transpiler: Formally Verifying Python Machine Learning Code

Machine learning model safety and reliability are receiving increasing attention, but verifying the correctness of Python machine learning code has long been a challenge. Recently, a project called "Provably Correct Tensor Puzzles" proposed formally verifying Python machine learning code by building a Jax-Lean transpiler. This progress, shared by community members on Twitter, demonstrates the idea of combining the Lean theorem prover with the Jax deep learning framework, offering a new path toward the reliability of machine learning code.

Why Do We Need Formal Verification of Machine Learning Code?

Python machine learning code is typically written in a dynamic, imperative style and lacks rigorous mathematical guarantees. Even if a model performs well on a test set, potential bugs or adversarial vulnerabilities cannot be ruled out. Formal verification uses mathematical proofs to ensure that code satisfies its specification under all possible inputs, but it has traditionally been confined to safety-critical domains. Bringing formal methods into machine learning means being able to prove the correctness of core logic such as tensor operations and gradient computation, thereby increasing model trustworthiness.

The core motivation of this work is: if Python machine learning code can be transpiled into a formal representation in Lean, then Lean's proving capabilities can be used to verify its properties. This is not merely academic exploration; it also provides a tooling foundation for deploying highly reliable AI systems in the future.

Formal verification is a method for proving software correctness based on mathematical logic. It requires developers to precisely describe the properties a program should satisfy using a formal specification language (such as Lean, Coq, or Isabelle), and then use a theorem prover to construct machine-checkable proofs ensuring that the program satisfies the specification under all legal inputs. Unlike traditional testing, formal verification does not rely on sampling coverage but exhausts all possible states, so it can find boundary errors that testing struggles to reach. The cost is a steep learning curve and high verification expense, so it is usually used only for safety-critical systems such as compilers, operating system kernels, and security protocols. In recent years, as tool ecosystems such as Lean have matured, formal methods have been gradually penetrating numerical computing and machine learning.

How Does the Jax-Lean Transpiler Work?

According to the original tweet, the author built a Jax-to-Lean transpiler that converts tensor computation graphs in Jax into expressions in Lean. Jax itself is functional, and its computation graphs are well suited to formal modeling. The transpiler likely maps Jax primitive operations (such as jnp.dot and jnp.sum) to tensor operation definitions in Lean and uses Lean's dependent type system to express shape constraints and numerical properties.

The transpilation process is not a simple syntax conversion; it requires establishing a correspondence between Jax semantics and Lean semantics. The author uses "tensor puzzles" as test cases. These puzzles usually involve complex indexing and broadcasting operations and are good benchmarks for verifying the correctness of the transpiler. Although the tweet did not reveal implementation details, this approach is consistent with the recent trend of using Lean to verify numerical algorithms.

Challenges of Formal Verification

Machine learning code involves floating-point arithmetic, randomness, and large-scale tensors, and full formal verification still faces challenges. Examples include modeling floating-point error and proving automatic differentiation rules. However, Jax's functional characteristics reduce side effects, making transpilation and verification more feasible. This work may focus on deterministic tensor operations as a step toward full verification.

Jax is a high-performance numerical computing library developed by Google, whose core features are the functional programming paradigm and automatic differentiation. Jax traces Python functions into purely functional computation graphs (jaxpr), in which each node is a side-effect-free primitive operation such as matrix multiplication, reduction, or broadcasting. This intermediate representation is naturally suited to formal modeling: the computation graph can be mapped node by node to expressions in Lean, and Lean's dependent type system can precisely express tensor shape and dimension constraints. Compared with PyTorch's imperative execution, Jax's purely functional characteristics eliminate the verification complexity brought by mutable state, making transpilation and proof more feasible.

Lean is an interactive theorem prover and also a dependently typed programming language. It allows users to define mathematical objects, state theorems, and construct machine-checkable proofs. Lean's mathematics library Mathlib already contains a large number of algebraic, analytical, and topological results, providing a foundation for verifying numerical algorithms. In Lean, tensors can be represented as indexed finite-dimensional array types, and shape constraints are statically checked at compile time through dependent types, thereby turning runtime errors into type errors.

Formalization of floating-point arithmetic is a core difficulty in numerical verification. IEEE 754 floating-point numbers do not satisfy associativity and distributivity in the real number field, rounding errors accumulate and propagate, and Lean's standard mathematics library is mainly based on exact real and rational numbers. To reason about the correctness of floating-point tensor operations in Lean, it is necessary to introduce floating-point error models, interval analysis, or error bound proofs, which significantly increases the verification burden. In addition, automatic differentiation involves backward traversal of computation graphs and chained composition of gradient rules, and its correctness proof requires establishing a refinement relation between differential semantics and numerical implementation. Current work very likely covers only deterministic tensor operations over integers or exact rational numbers, while floating point and randomness remain open problems to be tackled.

Significance for Machine Learning Engineering

If the Jax-Lean transpiler matures, developers can obtain formal guarantees while writing Jax code, reducing debugging and testing costs. Especially for safety-critical AI applications (such as healthcare and autonomous driving), this verification approach can provide stronger confidence. In addition, formal verification can help discover hidden boundary-case errors and improve code quality.

However, this is still an early exploration. The tweet only mentions that "a Jax-Lean transpiler was built" and does not yet specify its coverage or degree of automation. Community reaction has been positive, but practical application will have to wait for open sourcing and more case studies.

Conclusion and Outlook

"Provably Correct Tensor Puzzles" demonstrates the potential of the intersection between formal verification and machine learning engineering. Through the Jax-Lean transpiler, Python machine learning code may gain mathematical-level correctness guarantees. Although there is a long road ahead, this direction is worth watching. In the future, we may see more frameworks support formal verification, making AI systems more reliable.

One-line Summary

Jax-Lean transpiler converts Python machine learning code into formal representations in Lean to verify the correctness of tensor operations.