BurnpackStore

Cargo.toml

@@ -81,7 +81,8 @@ regex = { version = "1.11.3", default-features = false, features = [
 reqwest = { version = "0.12.23", default-features = false, features = [
     "rustls-tls",
 ] }
-rmp-serde = "1.3.0"
+ciborium = { version = "0.2", default-features = false }
+rmp-serde = { version = "1.3.0", default-features = false }
 rstest = "0.25.0"
 rusqlite = "0.37.0"
 rust-format = "0.3.4"

crates/burn-core/src/module/base.rs

@@ -1,4 +1,4 @@
-use super::{ParamId, Quantizer};
+use super::{Param, ParamId, Quantizer};
 use crate::{
     record::Record,
     tensor::backend::{AutodiffBackend, Backend},
@@ -19,11 +19,12 @@ macro_rules! module {
         impl<B: Backend> ModuleMapper<B> for Mapper {
             fn map_float<const D: usize>(
                 &mut self,
-                _id: ParamId,
-                tensor: Tensor<B, D>,
-            ) -> Tensor<B, D> {
+                param: Param<Tensor<B, D>>,
+            ) -> Param<Tensor<B, D>> {
+                let (id, tensor, mapper) = param.consume();
                 let func = $item;
-                func(tensor)
+                let tensor = func(tensor);
+                Param::from_mapped_value(id, tensor, mapper)
             }
         }
         let mut mapper = Mapper;
@@ -35,9 +36,9 @@ macro_rules! module {
             backend: core::marker::PhantomData<B>,
         }
         impl<'a, B: Backend> ModuleVisitor<B> for Visitor<'a, B> {
-            fn visit_float<const D: usize>(&mut self, _id: ParamId, tensor: &Tensor<B, D>) {
+            fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
                 let func = $item;
-                func(tensor, &mut self.state)
+                func(&param.val(), &mut self.state)
             }
         }
         #[allow(clippy::redundant_closure_call)]
@@ -211,29 +212,26 @@ pub trait Module<B: Backend>: Clone + Send + core::fmt::Debug {
 
 /// Module visitor trait for traversing and inspecting module parameters.
 pub trait ModuleVisitor<B: Backend> {
-    /// Visit a float tensor in the module.
+    /// Visit a float parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The float tensor to visit
+    /// - `param`: The float parameter to visit
     #[allow(unused_variables)]
-    fn visit_float<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D>) {}
+    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {}
 
-    /// Visit an int tensor in the module.
+    /// Visit an int parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The integer tensor to visit
+    /// - `param`: The integer parameter to visit
     #[allow(unused_variables)]
-    fn visit_int<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Int>) {}
+    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<B, D, Int>>) {}
 
-    /// Visit a bool tensor in the module.
+    /// Visit a bool parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The boolean tensor to visit
+    /// - `param`: The boolean parameter to visit
     #[allow(unused_variables)]
-    fn visit_bool<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Bool>) {}
+    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<B, D, Bool>>) {}
 
     /// Called when entering a submodule.
     ///
@@ -321,51 +319,49 @@ pub trait ModuleMapper<B: Backend> {
     #[allow(unused_variables)]
     fn exit_module(&mut self, name: &str, container_type: &str) {}
 
-    /// Map a float tensor in the module.
+    /// Map a float parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The float tensor to transform
+    /// - `param`: The float parameter to transform
     ///
     /// # Returns
-    /// The transformed tensor
+    /// The transformed parameter
     #[allow(unused_variables)]
-    fn map_float<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D>) -> Tensor<B, D> {
-        tensor
+    fn map_float<const D: usize>(&mut self, param: Param<Tensor<B, D>>) -> Param<Tensor<B, D>> {
+        let (id, tensor, mapper) = param.consume();
+        Param::from_mapped_value(id, tensor, mapper)
     }
 
-    /// Map an int tensor in the module.
+    /// Map an int parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The integer tensor to transform
+    /// - `param`: The integer parameter to transform
     ///
     /// # Returns
-    /// The transformed tensor
+    /// The transformed parameter
     #[allow(unused_variables)]
     fn map_int<const D: usize>(
         &mut self,
-        id: ParamId,
-        tensor: Tensor<B, D, Int>,
-    ) -> Tensor<B, D, Int> {
-        tensor
+        param: Param<Tensor<B, D, Int>>,
+    ) -> Param<Tensor<B, D, Int>> {
+        let (id, tensor, mapper) = param.consume();
+        Param::from_mapped_value(id, tensor, mapper)
     }
 
-    /// Map a bool tensor in the module.
+    /// Map a bool parameter in the module.
     ///
     /// # Parameters
-    /// - `id`: The unique identifier of the parameter
-    /// - `tensor`: The boolean tensor to transform
+    /// - `param`: The boolean parameter to transform
     ///
     /// # Returns
-    /// The transformed tensor
+    /// The transformed parameter
     #[allow(unused_variables)]
     fn map_bool<const D: usize>(
         &mut self,
-        id: ParamId,
-        tensor: Tensor<B, D, Bool>,
-    ) -> Tensor<B, D, Bool> {
-        tensor
+        param: Param<Tensor<B, D, Bool>>,
+    ) -> Param<Tensor<B, D, Bool>> {
+        let (id, tensor, mapper) = param.consume();
+        Param::from_mapped_value(id, tensor, mapper)
     }
 }
 

crates/burn-core/src/module/initializer.rs

@@ -107,6 +107,7 @@ impl Initializer {
         let device = device.clone();
         let shape: Shape = shape.into();
         let config = self.clone();
+        let shape_for_closure = shape.clone();
 
         Param::uninitialized(
             ParamId::new(),
@@ -123,6 +124,7 @@ impl Initializer {
             },
             device,
             true,
+            shape_for_closure,
         )
     }