Axiom.jl - Provably Correct Machine Learning
============================================
User API Layer
+-----------------------------------------------------------------+
| @axiom @ensure @prove* Sequential/Chain/Pipeline |
| Dense Conv2d BatchNorm LayerNorm MaxPool Dropout ReLU |
| CouplingLayer ActNorm Inv1x1Conv RevBlock NormFlow |
| train!() compile() verify() from_pytorch() |
+-----------------------------------------------------------------+
| | | |
v v v v
+----------+ +-----------+ +-----------+ +-----------+
| DSL & | | Training | | Verifi- | | Integra- |
| Macros | | Loop | | cation | | tions |
|----------| |-----------| |-----------| |-----------|
| axiom_ | | optimiz- | | proper- | | hugging- |
| macro | | ers.jl | | ties.jl | | face.jl |
| ensure | | loss.jl | | checker | | pytorch |
| prove* | | train.jl | | certifi- | | ext |
| pipeline | | gradient* | | cates.jl | | |
+----------+ +-----------+ | serial- | +-----------+
| ize.jl |
| proof_ |
| export |
+-----------+
|
+-----------------------------+-----------------------------+
| | |
v v v
+-----------+ +-----------+ +-----------+
| SMTLib.jl | | Lean 4 | | Coq / |
| (bundled) | | Export* | | Isabelle* |
|-----------| +-----------+ +-----------+
| z3, cvc5 |
| yices, |
| mathsat |
+-----------+
Backend Abstraction Layer (15 backends incl. SmartBackend)
+-----------------------------------------------------------------+
| abstract.jl - AbstractBackend / set_backend! / compile() |
| SmartBackend (per-op dispatch) / MixedPrecision / self-healing |
+-----------------------------------------------------------------+
| | |
v v v
+-----------+ +-----------+ +-----------+
| Julia | | Zig | | GPU |
| Backend | | Backend | | Backends |
|-----------| |-----------| |-----------|
| (default) | | zig/ | | CUDA |
| reference | | axiom | | ROCm |
| impl | | .zig | | Metal |
| | | SIMD+MT | | |
+-----------+ +-----------+ +-----------+
|
Coprocessor Backends (self-healing fallback) |
+-----------------------------------------------------------------+
| TPU | NPU | DSP | PPU | Math | FPGA | VPU | QPU | Crypto |
|-----------------------------------------------------------------|
| Environment-based detection (AXIOM_*_AVAILABLE) |
| Strict mode: AXIOM_*_REQUIRED=1 prevents fallback |
| Self-healing: graceful degradation to JuliaBackend |
+-----------------------------------------------------------------+
* = disabled, stub, or placeholder
| Component | Status | Progress |
|---|---|---|
Type System (Tensor) |
Done |
|
Layer Definitions |
Done |
|
Activation Functions |
Done |
|
Optimizers (SGD/Adam) |
Done |
|
Loss Functions |
Done |
|
Training Loop |
Done |
|
Model Save/Load |
Done |
|
Autograd (Zygote) |
Done |
|
Data Utilities |
Done |
|
Model Containers |
Done |
|
| Component | Status | Progress |
|---|---|---|
CouplingLayer (affine) |
Done |
|
ActNorm (data-dep init) |
Done |
|
Invertible1x1Conv (LU) |
Done |
|
RevBlock (reversible) |
Done |
|
InvertibleSequential |
Done |
|
NormalizingFlow |
Done |
|
Custom Zygote adjoints |
Done |
|
Roundtrip verification |
Done |
|
| Component | Status | Progress |
|---|---|---|
@axiom macro |
Done |
|
@ensure macro |
Done |
|
@prove macro |
Done |
|
Pipeline DSL |
Done |
|
| Component | Status | Progress |
|---|---|---|
Property Checking |
Done |
|
Proof Certificates |
Done |
|
Serialization |
Done |
|
SMTLib.jl (bundled) |
Done |
|
Lean 4 Export |
Done |
|
Coq Export |
Done |
|
Isabelle Export |
Done |
|
Proof Import |
Done |
|
| Component | Status | Progress |
|---|---|---|
matmul |
Done |
|
conv2d |
Done |
|
activations (all) |
Done |
|
batchnorm / layernorm |
Done |
|
pooling (max/avg/glob) |
Done |
|
dropout / flatten |
Done |
|
| Component | Status | Progress |
|---|---|---|
matmul (SIMD tiled) |
Done |
|
activations (15 funcs) |
Done |
|
conv2d + depthwise |
Done |
|
pooling (max/avg/glob) |
Done |
|
norm (batch/layer/rms) |
Done |
|
flash attention |
Done |
|
rotary embeddings |
Done |
|
Julia-side ccall wiring |
Done |
|
Compiled .so artifact |
Done |
|
FFI exports (32 syms) |
Done |
|
SIMD GELU/sigmoid/tanh |
Done |
|
Multi-threaded dispatch |
Done |
|
| Component | Status | Progress |
|---|---|---|
CUDA: matmul |
Done |
|
CUDA: relu |
Done |
|
CUDA: softmax |
Done |
|
CUDA: conv2d |
Done |
|
CUDA: batchnorm |
Done |
|
CUDA: pooling |
Done |
|
ROCm: matmul/relu/soft |
Done |
|
ROCm: conv2d/norm/pool |
Done |
|
Metal: matmul/relu/soft |
Done |
|
Metal: conv2d/norm/pool |
Done |
|
| Component | Status | Progress |
|---|---|---|
Backend type hierarchy |
Done |
|
Env-based detection |
Done |
|
Self-healing fallback |
Done |
|
Strict mode / required |
Done |
|
Runtime diagnostics |
Done |
|
Capability reporting |
Done |
|
TPU extension skeleton |
Skel |
|
NPU extension skeleton |
Skel |
|
DSP extension skeleton |
Skel |
|
PPU extension skeleton |
Skel |
|
Math extension skeleton |
Skel |
|
FPGA extension skeleton |
Skel |
|
VPU extension skeleton |
Skel |
|
QPU extension skeleton |
Skel |
|
Crypto ext. skeleton |
Skel |
|
| Component | Status | Progress |
|---|---|---|
PyTorch import/export |
Done |
|
HuggingFace framework |
Done |
|
HF model converters |
Done |
|
SafeTensors loader |
Done |
|
Model Metadata |
Done |
|
Resource-aware dispatch |
Done |
|
| Component | Status | Progress |
|---|---|---|
Backend dispatch |
Done |
|
SmartBackend (per-op) |
Done |
|
Mixed precision |
Done |
|
fold_batchnorm |
Done |
|
fold_constants |
Done |
|
dead code elimination |
Done |
|
aggressive opt pass |
Done |
|
| Component | Status | Progress |
|---|---|---|
CI/CD (21 workflows) |
Good |
|
Tests (283 passing) |
Good |
|
Benchmarks |
Done |
|
Documentation (wiki) |
Done |
|
RSR Compliance |
Done |
|
Bot directives (8) |
Done |
|
Contractiles (5) |
Done |
|
SCM files (5) |
Done |
|
| Dependency | Purpose | Required |
|---|---|---|
Julia >= 1.10 |
Runtime |
Yes |
LinearAlgebra |
Matrix operations |
Yes |
SHA |
Proof certificate hashing |
Yes |
JSON |
Metadata serialization |
Yes |
CUDA.jl |
NVIDIA GPU acceleration |
Optional |
AMDGPU.jl |
AMD GPU acceleration |
Optional |
Metal.jl |
Apple GPU acceleration |
Optional |
PyCall.jl |
PyTorch interop |
Optional |
Zig toolchain |
Zig backend compilation |
Optional |
Z3/CVC5 |
SMT solver for @prove |
Optional |
Strongest areas: Core layers, activations, reversible computing
(CouplingLayer/ActNorm/Inv1x1Conv/RevBlock/NormalizingFlow with custom
Zygote adjoints), Zig kernel implementations (sole native backend, 36
exports, 395KB .so, SIMD + 4-thread dispatch), SmartBackend per-op
dispatch, SIMD-optimized Zig kernels (GELU 3x, RMSNorm 7x, sigmoid
2.9x), multi-threaded element-wise ops (>64K threshold), SMTLib,
coprocessor dispatch infrastructure (9 backends with setup guides and
hardware detection), compile optimizations (incl. mixed precision with
loss scaling), certificates, HuggingFace (7 architectures
SafeTensors), RSR compliance (SPDX on all files), GPU extensions (full
coverage), autograd (Zygote), @prove (heuristic+SMT), proof export (real
tactics), backend-aware dispatch (LayerNorm/RMSNorm route through
backends), model save/load (binary + metadata bundle), external
benchmarks (Axiom vs Flux vs PyTorch), 283 tests passing Weakest
areas: Coprocessor skeletons (20% — need real hardware for end-to-end
integration)
Axiom.jl is the flagship package in julia-ecosystem (part of
developer-ecosystem monorepo). 13 sibling COMPUTE packages share the
same backend abstraction: Cladistics.jl, BowtieRisk.jl, Cliometrics.jl,
KnotTheory.jl, HackenbushGames.jl, QuantumCircuit.jl, SMTLib.jl,
PolyglotFormalisms.jl, Causals.jl, SiliconCore.jl, LowLevel.jl,
ZeroProb.jl, ProvenCrypto.jl.