Axiom.jl

Provably Correct Machine Learning — machine learning where bugs are caught at compile time, not in production.

Table of Contents

• What is Axiom.jl?
• Features
  • Compile-Time Shape Verification
  • Formal Verification
  • Model Interoperability
  • High Performance
  • Coprocessor Targets
  • Model Packaging + Registry Manifests
  • Verification Telemetry
  • Serving APIs
  • Interop APIs
• Quick Start
  • Installation
  • Hello World
  • With @axiom DSL
• Why Axiom.jl?
  • The Problem
  • The Solution
• Documentation
• Project Structure
• Roadmap
• Contributing
• Julia-First Verification
• License
• Acknowledgments

What is Axiom.jl?

Axiom.jl is a next-generation ML framework that combines:

• Compile-time verification — shape errors caught before runtime

• Formal guarantees — verification checks and certificate workflows

• Optional acceleration — Zig/GPU backend paths with explicit fallback behaviour

• Julia elegance — express models as mathematical specifications

using Axiom

model = Sequential(
    Dense(784, 128, relu),
    Dense(128, 10),
    Softmax()
)

x = Tensor(randn(Float32, 16, 784))
y = model(x)
result = verify(model, properties=[ValidProbabilities(), FiniteOutput()], data=[(x, nothing)])
@assert result.passed

Features

Compile-Time Shape Verification

# PyTorch: Runtime error after hours of training
# Axiom.jl: Compile error in milliseconds

@axiom BrokenModel begin
    input :: Tensor{Float32, (224, 224, 3)}
    features = input |> Conv(64, (3,3))
    output = features |> Dense(10)  # COMPILE ERROR!
    # "Shape mismatch: Conv output is (222,222,64), Dense expects vector"
end

Formal Verification

@axiom SafeClassifier begin
    # ...
    @ensure valid_probabilities(output)    # Runtime assertion
    @prove ∀x. sum(softmax(x)) == 1.0      # Experimental proof workflow
end

# Generate verification certificates
cert = verify(model) |> generate_certificate
save_certificate(cert, "fda_submission.cert")

Model Interoperability

# Import from a PyTorch checkpoint (.pt/.pth/.ckpt) via built-in Python bridge
# (requires python3 + torch in the selected runtime)
model = from_pytorch("model.pt")

# Or import canonical descriptor JSON
model2 = from_pytorch("model.pytorch.json")

# Export supported models to ONNX
to_onnx(model, "model.onnx", input_shape=(1, 3, 224, 224))

Current scope:

• from_pytorch(…): canonical descriptor import + direct .pt/.pth/.ckpt bridge.

• to_onnx(…): export for Sequential/Pipeline models built from Dense/Conv/Norm/Pool + common activations.

High Performance

# Development: Julia backend
model = Sequential(Dense(784, 128, relu), Dense(128, 10))

# Production path: optional Zig backend
prod_model = compile(model, backend=ZigBackend("/path/to/libaxiom_zig.so"), optimize=:aggressive)

Coprocessor Targets

Note

Status — skeleton, not accelerated execution. Coprocessor support today is a detection + dispatch surface: detect_coprocessor() probes for devices and compile(…, backend=cop) routes through the backend interface, but the per-device compute kernels are not yet implemented, so execution gracefully falls back to the Julia backend. Treat this as a forward-looking API, not hardware-accelerated inference.

# Non-GPU accelerator targets with self-healing fallback
cop = detect_coprocessor()  # TPU/NPU/VPU/QPU/PPU/MATH/CRYPTO/FPGA/DSP or nothing
if cop !== nothing
    model_accel = compile(model, backend=cop, verify=false)
end

Model Packaging + Registry Manifests

metadata = create_metadata(
    model;
    name="my-model",
    architecture="Sequential",
    version="1.0.0",
)
verify_and_claim!(metadata, "FiniteOutput", "verified=true; source=ci")

bundle = export_model_package(model, metadata, "build/model_package")
entry = build_registry_entry(bundle["manifest"]; channel="stable")
export_registry_entry(entry, "build/model_package/registry-entry.json")

Verification Telemetry

reset_verification_telemetry!()
result = verify(model, properties=[FiniteOutput()], data=[(x, nothing)])
run_payload = verification_result_telemetry(result; source="inference-gate")
summary = verification_telemetry_report()

Serving APIs

# REST
rest_server = serve_rest(model; host="0.0.0.0", port=8080, background=true)

# GraphQL
graphql_server = serve_graphql(model; host="0.0.0.0", port=8081, background=true)

# gRPC bridge server + contract generation
# - binary unary protobuf (`application/grpc`)
# - JSON bridge mode (`application/grpc+json`)
grpc_server = serve_grpc(model; host="0.0.0.0", port=50051, background=true)
generate_grpc_proto("axiom_inference.proto")

Interop APIs

# PyTorch import (checkpoint bridge or canonical descriptor JSON)
model = from_pytorch("model.pt")
model = from_pytorch("model.pytorch.json")

# ONNX export (Dense/Conv/Norm/Pool + common activations)
to_onnx(model, "model.onnx", input_shape=(1, 3, 224, 224))

Quick Start

Installation

using Pkg
Pkg.add("Axiom")

Hello World

using Axiom

# Define a simple classifier
model = Sequential(
    Dense(784, 256, relu),
    Dense(256, 10),
    Softmax()
)

# Generate sample data
x = randn(Float32, 32, 784)

# Inference
predictions = model(x)

# Verify properties
@ensure all(sum(predictions, dims=2) .≈ 1.0)

With @axiom DSL

using Axiom

@axiom MNISTClassifier begin
    input :: Tensor{Float32, (:batch, 28, 28, 1)}
    output :: Probabilities(10)

    features = input |> Conv(32, (3,3)) |> ReLU |> MaxPool((2,2))
    features = features |> Conv(64, (3,3)) |> ReLU |> MaxPool((2,2))
    flat = features |> GlobalAvgPool() |> Flatten
    output = flat |> Dense(64, 10) |> Softmax

    @ensure valid_probabilities(output)
end

model = MNISTClassifier()

Why Axiom.jl?

The Problem

ML models are deployed in critical applications:

• Medical diagnosis

• Autonomous vehicles

• Financial systems

• Criminal justice

Yet our tools allow bugs to slip through to production.

The Solution

Axiom.jl catches bugs before they cause harm:

Issue	PyTorch	Axiom.jl
Shape mismatch	Runtime crash	Compile error
NaN in output	Silent failure	Detected/proven
Invalid probabilities	Undetected	Checkable with verification properties
Adversarial fragility	Unknown	Roadmap / partial

Documentation

• Home — start here

• User Guide — install, infer, verify

• Developer Guide — build/test/release workflow

• Release Checklist — pre-release and release-day gates

• Vision — why we built this

• @axiom DSL — model definition guide

• Verification — @ensure and @prove

• Migration Guide — from PyTorch

• FAQ — common questions

• Testing Taxonomy — how the suite is classified (proven-tests categories + provenance tiers)

• Roadmap — tracked commitments and delivery criteria

Project Structure

Axiom.jl/
├── src/                 # Julia source
│   ├── Axiom.jl         # Main module
│   ├── types/           # Tensor type system
│   ├── layers/          # Neural network layers
│   ├── dsl/             # @axiom macro system
│   ├── verification/    # @ensure, @prove
│   ├── training/        # Optimizers, loss functions
│   └── backends/        # Backend abstraction (15 backends)
├── zig/                 # Zig native backend
│   └── src/             # matmul, conv, norm, attention, etc.
├── ext/                 # GPU package extensions (CUDA, ROCm, Metal)
├── test/                # Test suite
├── examples/            # Example models
└── docs/                # Documentation & wiki

Roadmap

• ✓ v0.1 — core framework, DSL, verification basics

• ❏ v0.2 — full Zig backend, GPU support

• ❏ v0.3 — HuggingFace integration, model zoo

• ❏ v0.4 — advanced proofs, SMT integration

• ❏ v1.0 — production ready, industry certifications

Contributing

We welcome contributions! See CONTRIBUTING.

• Bug reports and feature requests

• Documentation improvements

• New layers and operations

• Performance optimizations

• Verification methods

Julia-First Verification

Axiom’s proof system is Julia-native by default. SMT solving runs through packages/SMTLib.jl with no native backend dependency. The Zig SMT runner is an optional backend you can enable for hardened subprocess control.

Julia-native example:

@prove ∃x. x > 0

Optional Zig runner:

export AXIOM_SMT_RUNNER=zig
export AXIOM_ZIG_LIB=/path/to/libaxiom_zig.so
export AXIOM_SMT_SOLVER=z3

@prove ∃x. x > 0

License

Licensed under the Mozilla Public License 2.0 (MPL-2.0).

Acknowledgments

Axiom.jl builds on the shoulders of giants:

• Julia — the language

• Flux — inspiration for Julia ML

• Zig — native performance backend

• PyTorch — ecosystem compatibility

---

The future of ML is verified.

Get Started | Release Checklist | Roadmap | Star on GitHub