Move *towards* no-std support for bimm.

Author: crutcher

Cargo.lock

@@ -32,7 +32,7 @@ dirs = "^6.0.0"
 # > cudarc = { workspace = true, features = ["nvrtc"], optional = true }
 cudarc = "0.16.6"
 
-bimm-contracts = "^0.19.1"
+bimm-contracts = "^0.19.2"
 
 # Burn coupled-dependencies
 globwalk = "^0.9.1"

Cargo.toml

@@ -34,9 +34,8 @@ criterion = { workspace = true }
 
 
 [features]
-default = ["std"]
+default = []
 std = []
-wgpu = ["burn/wgpu"]
 
 [[bench]]
 name = "drop_block"

crates/bimm/Cargo.toml

@@ -1,6 +1,7 @@
 //! # Cache Policy
 
 use crate::cache::disk;
+use alloc::string::String;
 use anyhow::bail;
 use burn::config::Config;
 use burn::data::network::downloader;

crates/bimm/src/cache/disk.rs

@@ -3,11 +3,13 @@
 use crate::cache::weights::{
     PretrainedWeightsDescriptor, PretrainedWeightsMap, StaticPretrainedWeightsMap,
 };
+use alloc::collections::BTreeMap;
+use alloc::format;
+use alloc::string::{String, ToString};
+use alloc::sync::Arc;
 use anyhow::bail;
 use burn::config::Config;
-use std::collections::BTreeMap;
-use std::fmt::Debug;
-use std::sync::Arc;
+use core::fmt::Debug;
 
 /// Static builder for a [`PreFabConfig`]
 pub struct StaticPreFabConfig<C>
@@ -63,8 +65,8 @@ where
 {
     fn fmt(
         &self,
-        f: &mut std::fmt::Formatter<'_>,
-    ) -> std::fmt::Result {
+        f: &mut core::fmt::Formatter<'_>,
+    ) -> core::fmt::Result {
         self.to_prefab().fmt(f)
     }
 }
@@ -94,11 +96,11 @@ where
 {
     fn fmt(
         &self,
-        f: &mut std::fmt::Formatter<'_>,
-    ) -> std::fmt::Result {
+        f: &mut core::fmt::Formatter<'_>,
+    ) -> core::fmt::Result {
         let pretty = f.alternate();
 
-        let type_name = std::any::type_name::<C>();
+        let type_name = core::any::type_name::<C>();
         let mut handle = f.debug_struct(&format!("PreFabConfig<{}>", type_name));
 
         handle

crates/bimm/src/cache/prefabs.rs

@@ -1,9 +1,12 @@
 //! # Module / Weight Caches
 
 use crate::cache::disk::DiskCacheConfig;
+use alloc::collections::BTreeMap;
+use alloc::string::{String, ToString};
+use alloc::vec::Vec;
+use alloc::{format, vec};
 use anyhow::bail;
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
 use std::path::PathBuf;
 
 const X25: crc::Crc<u16> = crc::Crc::<u16>::new(&crc::CRC_16_IBM_SDLC);
@@ -197,6 +200,7 @@ impl PretrainedWeightsMap {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use alloc::string::{String, ToString};
 
     #[test]
     fn test_static_descriptor_to_descriptor() {

crates/bimm/src/cache/weights.rs

@@ -2,6 +2,8 @@
 //!
 //! Utilities for computing the output shape of convolution operations.
 
+use alloc::vec::Vec;
+
 /// Predict the output size of a 1D convolution operation.
 ///
 /// ```text

crates/bimm/src/compat/conv_shape.rs

@@ -1,4 +1,7 @@
 //! [`Tensor`] ops.
+
+use alloc::vec;
+use alloc::vec::Vec;
 use burn::prelude::{Backend, Tensor};
 use std::f64;
 

crates/bimm/src/compat/ops.rs

@@ -7,14 +7,15 @@ use crate::layers::drop::size_config::SizeConfig;
 use crate::utility::burn::kernels;
 use crate::utility::burn::noise::NoiseConfig;
 use crate::utility::probability::expect_probability;
+use alloc::format;
 use bimm_contracts::unpack_shape_contract;
 use burn::config::Config;
 use burn::module::{Content, Module, ModuleDisplay, ModuleDisplayDefault};
 use burn::prelude::{Backend, Float, Tensor};
 use burn::tensor::module::max_pool2d;
 use burn::tensor::{DType, Distribution};
+use core::ops::Range;
 use serde::{Deserialize, Serialize};
-use std::ops::Range;
 
 /// Configuration for `DropBlock`.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -224,7 +225,7 @@ impl DropBlockOptions {
         let [kh, kw] = self.kernel;
         let kh = kh.resolve(h);
         let kw = kw.resolve(w);
-        [std::cmp::min(h, kh), std::cmp::min(w, kw)]
+        [core::cmp::min(h, kh), core::cmp::min(w, kw)]
     }
 
     /// Compute the adjusted gamma rate.

crates/bimm/src/layers/drop/drop_block.rs

@@ -1,5 +1,7 @@
 //! Common rate table for `DropPath` regularization.
+
 use crate::compat::ops::float_vec_linspace;
+use alloc::vec::Vec;
 
 /// Computes a progressive incremental path drop rate for stochastic depth.
 ///
@@ -139,6 +141,7 @@ mod tests {
     use super::*;
     use crate::layers::drop::rate_table::progressive_dpr;
     use crate::testing::assert_close_to_vec;
+    use alloc::vec;
     use hamcrest::prelude::*;
 
     #[test]