- What It Actually Does Today
- Claim: Neural network layers with correct forward passes
- Claim: Compile-time shape verification via
@axiomDSL - Claim: Verification properties, certificates, and proof assistant export
- Claim: REST, GraphQL, and gRPC serving APIs
- Claim: PyTorch import and ONNX export
- Claim: Zig backend with optional SIMD / multi-threading
- How It Works
- Dogfooded Across The Account
- File Map
- Questions?
The README makes claims. This file backs them up with file paths, honest caveats, and a precise description of the critical paths through a sizeable codebase.
Axiom.jl is a next-generation ML framework that combines compile-time shape verification, formal property guarantees, optional Zig/GPU backend acceleration, and Julia elegance. Machine learning where bugs are caught at compile time, not in production. Shape errors caught before runtime. Verification checks and certificate workflows. Model packaging with registry manifests. REST, GraphQL, and gRPC serving. PyTorch import and ONNX export.
True and tested. Dense(in, out, activation) in src/layers/dense.jl supports
forward pass, optional bias, activation functions, and DimensionMismatch on wrong
input shapes. Conv2d(in_ch, out_ch, kernel_size; padding, stride) in
src/layers/conv.jl computes correct output spatial dimensions — the test verifies
a (4,32,32,3) input through a (3,64,(3,3)) conv produces (4,30,30,64), and (4,32,32,64)
with padding=1. Sequential, Chain, Residual containers are present.
Activation functions: relu, sigmoid, tanh, softmax, gelu, leaky_relu
and their in-place backend variants. BatchNorm, LayerNorm, Dropout,
MaxPool2d, AvgPool2d, GlobalAvgPool, Flatten are all exported.
Invertible layers (CouplingLayer, ActNorm, Invertible1x1Conv, RevBlock,
NormalizingFlow) with inverse, log_abs_det_jacobian, and forward_and_log_det.
DSL implemented; proof backend is partial. src/dsl/axiom_macro.jl implements
the @axiom ModelName begin … end macro. src/dsl/ensure.jl implements @ensure
for runtime assertion. src/dsl/prove.jl implements @prove — Julia-native by
default via packages/SMTLib.jl; an optional Zig SMT runner is triggered by the
AXIOM_SMT_RUNNER=zig environment variable and AXIOM_ZIG_LIB path. The
@prove ∃x. x > 0 example from the README works with the Julia-native backend.
Honest caveat: The @axiom macro’s shape-mismatch detection at "compile time"
means macro-expansion time, not Julia’s actual compile pass. For the Conv → Dense
mismatch example in the README, the error is raised when the model is constructed,
not when the file is parsed. This is earlier than a runtime crash but is not
type-level (dependent type) verification. The README’s PyTorch comparison is
directionally accurate.
src/verification/properties.jl defines ValidProbabilities, FiniteOutput,
NoNaN, NoInf as checkable property types. src/verification/checker.jl
implements verify(model, properties, data) returning VerificationResult with
passed::Bool. src/verification/certificates.jl provides ProofCertificate,
generate_certificate, save_certificate, load_certificate. Telemetry:
reset_verification_telemetry!, verification_result_telemetry,
verification_telemetry_report. Proof assistant export (src/proof_export.jl):
export_lean, export_coq, export_isabelle, proof_obligation_manifest,
reconcile_proof_bundle.
src/serving/api.jl implements serve_rest(model; host, port, background),
serve_graphql(model; …), serve_grpc(model; …) and generate_grpc_proto.
HTTP is a hard dependency in Project.toml. Tests verify the serving functions
exist and that the gRPC proto generation produces a valid .proto file.
src/integrations/interop.jl exports from_pytorch and to_onnx. from_pytorch
supports both direct .pt/.pth/.ckpt checkpoint import (via optional PyCall)
and canonical descriptor JSON import. to_onnx supports Sequential/Pipeline
models with Dense/Conv/Norm/Pool layers and common activations. The PyCall path
is declared as a weak dependency extension (AxiomPyTorchExt).
src/backends/zig_ffi.jl defines ZigBackend and init_zig_backend(lib_path).
The init function in src/Axiom.jl loads the Zig shared library when
AXIOM_ZIG_LIB is set in the environment, with a graceful warning fallback.
src/zig/ contains the Zig source for matmul, conv, norm, attention, and related
kernels. The Zig backend is also used for the optional SMT runner path in @prove.
Module load sequence (critical path order in src/Axiom.jl):
-
Types:
src/types/tensor.jl(Tensor,DynamicTensor,Shape),src/types/shapes.jl. -
Layers:
abstract.jl→dense.jl→conv.jl→activations.jl→normalization.jl→pooling.jl. -
DSL macros:
axiom_macro.jl→ensure.jl→prove.jl→pipeline.jl. -
Autograd:
gradient.jl→tape.jl→invertible.jl→invertible_rules.jl. -
Training:
optimizers.jl→loss.jl→train.jl. -
Verification:
properties.jl→checker.jl→certificates.jl→serialization.jl→proof_export.jl. -
Backends:
abstract.jl→julia_backend.jl→gpu_hooks.jl→zig_ffi.jl. -
Model packaging:
model_metadata.jl→model_packaging.jl. -
Utilities and serving:
utils/→serving/api.jl→integrations/interop.jl.
GPU backends (CUDA, ROCm, Metal) and coprocessor backends (TPU, NPU, DSP, PPU,
Math, FPGA, VPU, QPU, Crypto) are declared as [extensions] in Project.toml
and loaded conditionally. The PyTorch bridge is also an extension (AxiomPyTorchExt
requiring PyCall).
| Project | Connection |
|---|---|
|
ZeroProb types are directly usable in Axiom |
|
|
|
Vendored ( |
|
Shares the backend dispatch pattern (JuliaBackend → hardware fallthrough). QuantumCircuit was an earlier proving ground for the architecture. |
|
Model metadata and registry entry JSON is designed to be stored in per-project VeriSimDB instances for queryable model provenance. |
|
Axiom’s invertible layers and |
|
|
| Path | What’s There |
|---|---|
|
Top-level module (~250 lines): all includes in dependency order; ~150 exports covering every public type, function, macro, and backend symbol. |
|
|
|
|
|
|
>` chain sugar for model definition. |
|
|
|
|
|
|
|
Lean/Coq/Isabelle export, proof obligation manifest, bundle reconciliation, certificate import from proof assistants. |
|
|
|
|
|
|
|
|
|
|
|
Zig native backend source: matmul, conv, normalization, attention, and related compute kernels. |
|
GPU extensions (CUDA, ROCm, Metal) and 9 coprocessor extensions plus PyTorch bridge (PyCall). |
|
Comprehensive suite: tensor creation, Dense/Conv2d forward passes and error
cases, Sequential model, activation functions, BatchNorm/Dropout, verify with
ValidProbabilities/FiniteOutput, |
|
Internal Julia-native SMT library consumed by the |
|
Package identity (uuid |
|