Expand description
AxonML JIT — Tracing, IR, Optimization, and Cranelift Codegen
Top-level crate module for the AxonML just-in-time compiler. Re-exports the
tracing surface (Tracer, TracedValue, trace), typed IR (Graph,
Node, NodeId, Op, Shape, DataType), six-pass optimizer
(Optimizer, OptimizationPass covering constant folding, DCE, CSE,
algebraic simplification, elementwise fusion, and strength reduction), the
Cranelift-backed code generator (JitCompiler, CompiledFunction), the
higher-level compile_fn / compile_graph / CompiledModel /
LazyCompiled facade with CompileConfig, CompileStats, Backend, and
Mode, the FunctionCache LRU for compiled-function reuse, and the error
types JitError / JitResult. Unit tests exercise a basic trace round-trip
and a constant-folding optimization pass.
§File
crates/axonml-jit/src/lib.rs
§Author
Andrew Jewell Sr. — AutomataNexus LLC ORCID: 0009-0005-2158-7060
§Updated
April 16, 2026 11:15 PM EST
§Disclaimer
Use at own risk. This software is provided “as is”, without warranty of any kind, express or implied. The author and AutomataNexus shall not be held liable for any damages arising from the use of this software.
Re-exports§
pub use cache::FunctionCache;pub use codegen::CompiledFunction;pub use codegen::JitCompiler;pub use compile::Backend;pub use compile::CompileConfig;pub use compile::CompileStats;pub use compile::CompiledModel;pub use compile::LazyCompiled;pub use compile::Mode;pub use compile::compile_fn;pub use compile::compile_fn_with_config;pub use compile::compile_graph;pub use compile::compile_graph_with_config;pub use error::JitError;pub use error::JitResult;pub use ir::DataType;pub use ir::Graph;pub use ir::Node;pub use ir::NodeId;pub use ir::Op;pub use ir::Shape;pub use optimize::OptimizationPass;pub use optimize::Optimizer;pub use trace::TracedValue;pub use trace::Tracer;pub use trace::trace;
Modules§
- cache
- Function Cache — Hash-Keyed Reuse of Compiled JIT Graphs
- codegen
- Code Generation — Interpreter and Cranelift Native Backends
- compile
- torch.compile Equivalent - High-Level Compilation API
- error
- JIT Error Types — Compilation and Runtime Failures
- ir
- Intermediate Representation — Typed Tensor Graph for JIT
- optimize
- Graph Optimization — Six-Pass IR Transformation Pipeline
- trace
- Operation Tracing — Thread-Local Graph Recorder