Expand description
Reverse-mode automatic differentiation engine for AxonML.
Variable wraps Tensor<f32> with gradient tracking. backward() walks
the dynamic computation graph and accumulates gradients via GradFn /
GradientFunction traits. Gradient functions cover arithmetic, activations,
matmul, conv, RNN, loss, and attention ops. NoGradGuard disables tracking;
checkpoint / checkpoint_sequential trade compute for memory; amp
provides F16 autocast for mixed-precision training. The inspect module
offers trace_backward, to_dot (DOT graph export), node_count, depth,
and gradient flow summary — native graph visualization without external tools.
Modules: variable, backward, grad_fn, functions (basic, activation,
linalg, loss, conv, rnn, attention), graph, inspect, checkpoint,
amp, no_grad.
§File
crates/axonml-autograd/src/lib.rs
§Author
Andrew Jewell Sr. — AutomataNexus LLC ORCID: 0009-0005-2158-7060
§Updated
April 14, 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 amp::AutocastGuard;pub use amp::AutocastPolicy;pub use amp::autocast;pub use amp::autocast_dtype;pub use amp::disable_autocast;pub use amp::is_autocast_enabled;pub use backward::backward;pub use checkpoint::checkpoint;pub use checkpoint::checkpoint_rng_seed;pub use checkpoint::checkpoint_sequential;pub use functions::FusedAttentionBackward;pub use functions::GruGatesBackward;pub use functions::IdentityBackward;pub use functions::LstmGatesBackward;pub use grad_fn::GradFn;pub use grad_fn::GradientFunction;pub use graph::ComputationGraph;pub use graph::GraphNode;pub use inspect::GraphSnapshot;pub use inspect::SnapshotNode;pub use inspect::depth;pub use inspect::node_count;pub use inspect::to_dot;pub use inspect::trace_backward;pub use no_grad::NoGradGuard;pub use no_grad::no_grad;pub use variable::Variable;
Modules§
- amp
- Automatic Mixed Precision (AMP) — F16 autocast for faster training.
- backward
- Backward pass — reverse-mode gradient computation.
- checkpoint
- Gradient checkpointing — trade compute for memory during training.
- functions
- Gradient function implementations for all differentiable operations.
- grad_fn
- Gradient function traits — the differentiable operation interface.
- graph
- Dynamic computation graph construction for automatic differentiation.
- inspect
- Computation graph inspection and visualization — a novel AxonML feature.
- no_grad
- No-grad context — disable gradient tracking for inference / eval.
- prelude
- Convenient imports for common autograd usage.
- variable
Variable— tensor with automatic gradient tracking.