burn-autodiff: support no_std

Cargo.toml

@@ -167,5 +167,7 @@ cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features
 ### For xtask crate ###
 tracel-xtask = { version = "=1.1.8" }
 
+cfg_block = "0.2.0"
+
 [profile.dev]
 debug = 0 # Speed up compilation time and not necessary.

crates/burn-autodiff/Cargo.toml

@@ -25,6 +25,9 @@ burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "0.17.0", opt
 derive-new = { workspace = true }
 spin = { workspace = true }
 log = { workspace = true }
+cfg_block = { workspace = true }
+hashbrown = { workspace = true }
+num-traits = { workspace = true }
 
 [dev-dependencies]
 burn-tensor = { path = "../burn-tensor", version = "0.17.0", default-features = false, features = [

crates/burn-autodiff/src/backend.rs

@@ -4,6 +4,7 @@ use crate::{
     runtime::AutodiffClient,
     tensor::AutodiffTensor,
 };
+use alloc::{format, string::String};
 use burn_tensor::{
     backend::{AutodiffBackend, Backend},
     ops::{BoolTensor, IntTensor, QuantizedTensor},

crates/burn-autodiff/src/checkpoint/base.rs

@@ -2,8 +2,10 @@ use super::{
     retro_forward::RetroForwards,
     state::{BackwardStates, State},
 };
+use crate::collections::HashMap;
 use crate::graph::NodeID;
-use std::collections::HashMap;
+
+use alloc::{vec, vec::Vec};
 
 #[derive(new, Debug)]
 /// Links a [NodeID] to its autodiff graph [NodeRef]

crates/burn-autodiff/src/checkpoint/builder.rs

@@ -1,9 +1,11 @@
 use crate::{
+    collections::HashMap,
     graph::{ComputingProperty, NodeID, NodeSteps},
     tensor::AutodiffTensor,
 };
+use alloc::{boxed::Box, sync::Arc, vec::Vec};
 use burn_tensor::backend::Backend;
-use std::{any::Any, collections::HashMap, sync::Arc};
+use core::any::Any;
 
 use super::{
     base::{Checkpointer, NodeTree},

crates/burn-autodiff/src/checkpoint/retro_forward.rs

@@ -1,6 +1,8 @@
+use crate::collections::HashMap;
 use crate::graph::NodeID;
 
-use std::{collections::HashMap, fmt::Debug, sync::Arc};
+use alloc::sync::Arc;
+use core::fmt::Debug;
 
 use super::state::{BackwardStates, State};
 

crates/burn-autodiff/src/checkpoint/state.rs

@@ -1,6 +1,8 @@
-use std::{any::Any, collections::HashMap};
+use core::any::Any;
 
+use crate::collections::HashMap;
 use crate::graph::NodeID;
+use alloc::boxed::Box;
 
 /// In order to accept arbitrary node output in the same hashmap, we need to upcast them to any.
 pub(crate) type StateContent = Box<dyn Any + Send>;

crates/burn-autodiff/src/checkpoint/strategy.rs

@@ -1,9 +1,9 @@
 use core::fmt::Debug;
-use std::sync::Arc;
 
 use burn_tensor::backend::Backend;
 
 use crate::{graph::ComputingProperty, tensor::AutodiffTensor};
+use alloc::sync::Arc;
 
 use super::{
     builder::{ActionType, CheckpointerBuilder},

crates/burn-autodiff/src/graph/base.rs

@@ -1,9 +1,9 @@
 use super::NodeID;
-use crate::{checkpoint::base::Checkpointer, grads::Gradients};
-use std::collections::HashMap;
+use crate::{checkpoint::base::Checkpointer, collections::HashMap, grads::Gradients};
+use alloc::{boxed::Box, vec::Vec};
 
 /// Backward step for reverse mode autodiff.
-pub trait Step: Send + std::fmt::Debug {
+pub trait Step: Send + core::fmt::Debug {
     /// Executes the step and consumes it.
     fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer);
     /// Depth of the operation relative to the first node added to a graph.

crates/burn-autodiff/src/graph/node.rs

@@ -1,5 +1,5 @@
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::Arc;
+use alloc::{sync::Arc, vec::Vec};
+use core::sync::atomic::{AtomicU64, Ordering};
 
 use crate::checkpoint::retro_forward::RetroForward;
 use crate::runtime::AutodiffClientImpl;

crates/burn-autodiff/src/graph/traversal.rs

@@ -1,6 +1,9 @@
 use super::{Step, StepBoxed};
-use crate::NodeID;
-use std::collections::{HashMap, HashSet};
+use crate::{
+    collections::{HashMap, HashSet},
+    NodeID,
+};
+use alloc::vec::Vec;
 
 /// Breadth for search algorithm.
 pub struct BreadthFirstSearch;

crates/burn-autodiff/src/lib.rs

@@ -1,3 +1,4 @@
+#![cfg_attr(not(feature = "std"), no_std)]
 #![warn(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 
@@ -34,3 +35,15 @@ pub use backend::*;
 
 #[cfg(feature = "export_tests")]
 mod tests;
+
+/// A facade around for HashMap and HashSet.
+/// This avoids elaborate import wrangling having to happen in every module.
+mod collections {
+    cfg_block::cfg_block! {
+        if #[cfg(feature = "std")] {
+            pub use std::collections::{HashMap, HashSet};
+        } else {
+            pub use hashbrown::{HashMap, HashSet};
+        }
+    }
+}

crates/burn-autodiff/src/ops/activation.rs

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use core::marker::PhantomData;
 
 use crate::{
     checkpoint::{

crates/burn-autodiff/src/ops/backward.rs

@@ -14,13 +14,13 @@ use burn_tensor::backend::Backend;
 /// Concrete types implementing this trait should not have any state.
 /// If a state is necessary during the backward pass,
 /// they should be declared with the associated type 'State'.
-pub trait Backward<B, const N: usize>: Send + std::fmt::Debug
+pub trait Backward<B, const N: usize>: Send + core::fmt::Debug
 where
     Self: Sized + 'static,
     B: Backend,
 {
     /// Associated type to compute the backward pass.
-    type State: Clone + Send + std::fmt::Debug + 'static;
+    type State: Clone + Send + core::fmt::Debug + 'static;
 
     /// The backward pass.
     fn backward(

crates/burn-autodiff/src/ops/base.rs

@@ -10,8 +10,9 @@ use crate::{
     graph::{ComputingProperty, NodeID, NodeRef, Requirement, Step},
     tensor::AutodiffTensor,
 };
+use alloc::{boxed::Box, vec::Vec};
 use burn_tensor::{backend::Backend, ops::FloatTensor, Shape, TensorMetadata};
-use std::marker::PhantomData;
+use core::marker::PhantomData;
 
 /// Operation in preparation.
 ///
@@ -134,7 +135,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Prepare an operation that requires a state during the backward pass.
@@ -161,7 +162,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, UnTracked>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Finish the preparation of an untracked operation and returns the output tensor.
@@ -184,7 +185,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Tracked>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Finish the preparation of a tracked operation and returns the output tensor.
@@ -235,7 +236,7 @@ struct OpsStep<B, T, SB, const N: usize>
 where
     B: Backend,
     T: Backward<B, N, State = SB>,
-    SB: Clone + Send + std::fmt::Debug + 'static,
+    SB: Clone + Send + core::fmt::Debug + 'static,
 {
     ops: Ops<SB, N>,
     backward: T,
@@ -246,7 +247,7 @@ impl<B, T, SB, const N: usize> Step for OpsStep<B, T, SB, N>
 where
     B: Backend,
     T: Backward<B, N, State = SB>,
-    SB: Clone + Send + std::fmt::Debug + 'static,
+    SB: Clone + Send + core::fmt::Debug + 'static,
 {
     fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer) {
         self.backward.backward(self.ops, grads, checkpointer);

crates/burn-autodiff/src/ops/bool_tensor.rs

@@ -1,4 +1,5 @@
 use crate::{checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor, Autodiff};
+use alloc::vec::Vec;
 
 use burn_tensor::{
     backend::Backend,
@@ -31,7 +32,7 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
         B::bool_reshape(tensor, shape)
     }
 
-    fn bool_slice(tensor: BoolTensor<B>, ranges: &[std::ops::Range<usize>]) -> BoolTensor<B> {
+    fn bool_slice(tensor: BoolTensor<B>, ranges: &[core::ops::Range<usize>]) -> BoolTensor<B> {
         B::bool_slice(tensor, ranges)
     }
 
@@ -41,7 +42,7 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
 
     fn bool_slice_assign(
         tensor: BoolTensor<Self>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
         value: BoolTensor<Self>,
     ) -> BoolTensor<Self> {
         B::bool_slice_assign(tensor, ranges, value)

crates/burn-autodiff/src/ops/int_tensor.rs

@@ -1,4 +1,5 @@
 use crate::{checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor, Autodiff};
+use alloc::vec::Vec;
 
 use burn_tensor::{
     backend::Backend,
@@ -27,7 +28,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_reshape(tensor, shape)
     }
 
-    fn int_slice(tensor: IntTensor<B>, ranges: &[std::ops::Range<usize>]) -> IntTensor<B> {
+    fn int_slice(tensor: IntTensor<B>, ranges: &[core::ops::Range<usize>]) -> IntTensor<B> {
         B::int_slice(tensor, ranges)
     }
 
@@ -37,7 +38,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
 
     fn int_slice_assign(
         tensor: IntTensor<B>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
         value: IntTensor<B>,
     ) -> IntTensor<B> {
         B::int_slice_assign(tensor, ranges, value)
@@ -305,7 +306,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_random(shape, distribution, device)
     }
 
-    fn int_arange(range: std::ops::Range<i64>, device: &Device<Self>) -> IntTensor<Self> {
+    fn int_arange(range: core::ops::Range<i64>, device: &Device<Self>) -> IntTensor<Self> {
         B::int_arange(range, device)
     }
 

crates/burn-autodiff/src/ops/qtensor.rs

@@ -1,4 +1,4 @@
-use std::ops::Range;
+use core::ops::Range;
 
 use burn_tensor::{
     backend::Backend,