Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

History

126 Commits

Axiom.jl

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

The @axiom macro tracks tensor shapes through the layer chain at macro-expansion time, so a provable mismatch is a compile-time error — raised before the model is ever constructed or run, not after hours of training.

# Correct: the feature dimensions chain cleanly (784 -> 256 -> 10).@axiom Classifier begin
input ::Tensor{Float32, (:batch, 784)}
output ::Tensor{Float32, (:batch, 10)}
hidden = input |>Dense(784, 256, relu)
output = hidden |>Dense(256, 10) |> Softmax
end# Broken: a Dense expecting 128 features fed a 256-feature tensor.@axiom BrokenModel begin
input ::Tensor{Float32, (:batch, 784)}
output ::Tensor{Float32, (:batch, 10)}
hidden = input |>Dense(784, 256, relu)
output = hidden |>Dense(128, 10) # COMPILE ERRORend# ERROR: @axiom BrokenModel: compile-time shape mismatch in `output`.# Dense layer expects 128 input feature(s), but the incoming tensor has 256.# Running shape entering this layer: (:batch, 256).

Verification is sound but incomplete: it rejects any mismatch it can prove from the declared input ::/output :: shapes and literal layer sizes (the Dense feature dimensions along the chain, and the final output declaration), and passes through anything it cannot statically resolve — Conv/pooling output geometry, non-literal layer arguments — so a valid model never receives a false compile error.

Formal Verification

@axiom SafeClassifier begin# ...@ensurevalid_probabilities(output) # Runtime assertion@prove ∀x. sum(softmax(x)) ==1.0# Experimental proof workflowend# Generate verification certificates
cert =verify(model) |> generate_certificate
save_certificate(cert, "fda_submission.cert")

Model Interoperability

# Import from a canonical PyTorch descriptor JSON — pure Julia, no runtime deps.
model =from_pytorch("model.pytorch.json")
# Raw .pt/.pth/.ckpt files are Python-pickle and need a PyTorch/Python runtime# to read, so they are not imported directly — export your model to the# `axiom.pytorch.sequential.v1` descriptor first, then import the .json above.# Export supported models to ONNXto_onnx(model, "model.onnx", input_shape=(1, 3, 224, 224))

Current scope:

  • from_pytorch(…): pure-Julia axiom.pytorch.sequential.v1 descriptor import.

  • 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 nothingif 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 — pure-Julia descriptor JSON (export from PyTorch first)
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@ensureall(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
@ensurevalid_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:

IssuePyTorchAxiom.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

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

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.

About

Julia package: Axiom

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

2 stars

Watchers

0 watching

Forks

Releases

Sponsor this project

Packages

Used by

Contributors

Languages