burn_autodiff/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! # Burn Autodiff
6//!
7//! This autodiff library is a part of the Burn project. It is a standalone crate
8//! that can be used to perform automatic differentiation on tensors. It is
9//! designed to be used with the Burn Tensor crate, but it can be used with any
10//! tensor library that implements the `Backend` trait.
11
12#[macro_use]
13extern crate derive_new;
14
15extern crate alloc;
16
17/// Checkpoint module.
18pub mod checkpoint;
19/// Gradients module.
20pub mod grads;
21/// Operation module.
22pub mod ops;
23
24pub(crate) mod graph;
25// Exported for backend extension
26pub use graph::NodeId;
27pub(crate) mod tensor;
28pub(crate) mod utils;
29
30mod backend;
31
32pub(crate) mod runtime;
33
34pub use backend::*;
35
36#[cfg(feature = "export_tests")]
37mod tests;
38
39/// A facade around for HashMap and HashSet.
40/// This avoids elaborate import wrangling having to happen in every module.
41mod collections {
42    #[cfg(not(feature = "std"))]
43    pub use hashbrown::{HashMap, HashSet};
44    #[cfg(feature = "std")]
45    pub use std::collections::{HashMap, HashSet};
46}