burn_core/module/param/
tensor.rs

1use super::reparameterization_dyn::{self, DynReparameterization};
2use super::{Param, ParamId, Parameter, Reparameterization};
3use crate::module::{
4    AutodiffModule, Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper,
5    ModuleVisitor,
6};
7use alloc::{boxed::Box, format, string::ToString, vec::Vec};
8use burn_tensor::{Bool, Device, Float, Int, Tensor, TensorData};
9
10impl<const D: usize> super::sealed::Sealed for Tensor<D, Float> {
11    fn materialize(self, reparameterization: &dyn DynReparameterization) -> Self {
12        *reparameterization
13            .materialize_dyn(Box::new(self))
14            .downcast::<Tensor<D>>()
15            .expect("Reparameterization should preserve tensor rank")
16    }
17}
18impl<const D: usize> super::sealed::Sealed for Tensor<D, Int> {}
19impl<const D: usize> super::sealed::Sealed for Tensor<D, Bool> {}
20
21impl<const D: usize> Parameter for Tensor<D, Float> {
22    fn device(&self) -> Device {
23        Tensor::device(self)
24    }
25
26    fn is_require_grad(&self) -> bool {
27        Tensor::is_require_grad(self)
28    }
29
30    fn set_require_grad(self, require_grad: bool) -> Self {
31        Tensor::set_require_grad(self, require_grad)
32    }
33
34    fn shape(&self) -> burn_std::Shape {
35        Tensor::shape(self)
36    }
37
38    fn load_to_device(self, device: &Device) -> Self {
39        if self.device() != *device {
40            Tensor::to_device(self, device).detach()
41        } else {
42            self
43        }
44    }
45}
46
47impl<const D: usize> Parameter for Tensor<D, Int> {
48    fn device(&self) -> Device {
49        Tensor::device(self)
50    }
51
52    fn is_require_grad(&self) -> bool {
53        false
54    }
55
56    fn set_require_grad(self, _require_grad: bool) -> Self {
57        self
58    }
59
60    fn shape(&self) -> burn_std::Shape {
61        Tensor::shape(self)
62    }
63
64    fn load_to_device(self, device: &Device) -> Self {
65        if self.device() != *device {
66            Tensor::to_device(self, device)
67        } else {
68            self
69        }
70    }
71}
72
73impl<const D: usize> Parameter for Tensor<D, Bool> {
74    fn device(&self) -> Device {
75        Tensor::device(self)
76    }
77
78    fn is_require_grad(&self) -> bool {
79        false
80    }
81
82    fn set_require_grad(self, _require_grad: bool) -> Self {
83        self
84    }
85
86    fn shape(&self) -> burn_std::Shape {
87        Tensor::shape(self)
88    }
89
90    fn load_to_device(self, device: &Device) -> Self {
91        if self.device() != *device {
92            Tensor::to_device(self, device)
93        } else {
94            self
95        }
96    }
97}
98
99impl<const D: usize> Param<Tensor<D>> {
100    /// Create a new parameter from a float tensor.
101    ///
102    /// # Warnings
103    ///
104    /// We strongly recommend using [Param::uninitialized] if you are using this method to
105    /// initialize parameters inside a module, since the tensor initialization will be lazy,
106    /// making the loading of weights more performant.
107    pub fn from_tensor(value: Tensor<D>) -> Self {
108        // When creating a parameter from a float tensor, we automatically mark it as requiring
109        // gradients, so that it can be updated by an optimizer.
110        Param::initialized(ParamId::new(), value.require_grad())
111    }
112
113    /// Create a new parameter from data.
114    pub fn from_data<T>(data: T, device: &Device) -> Self
115    where
116        T: Into<TensorData>,
117    {
118        let data: TensorData = data.into();
119        // When creating a parameter from a float tensor, we automatically mark it as requiring
120        // gradients, so that it can be updated by an optimizer.
121        device.memory_persistent_allocations(data, |data| {
122            let value = Tensor::from_data(data, device);
123            Param::initialized(ParamId::new(), value.require_grad())
124        })
125    }
126
127    /// Attach a custom or built-in reparameterization, replacing any existing one.
128    pub(crate) fn with_reparameterization<R>(mut self, reparameterization: R) -> Self
129    where
130        R: Reparameterization,
131    {
132        self.reparameterization = Some(reparameterization_dyn::boxed::<R, D>(reparameterization));
133        self
134    }
135}
136
137impl<const D: usize> Module for Param<Tensor<D>> {
138    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
139        match self.reparameterization_dyn() {
140            None => visitor.visit_float(self),
141            Some(reparameterization) => {
142                visitor.visit_float(&self.without_reparameterization());
143                visitor.enter_module(reparameterization.name(), "Reparameterization");
144                reparameterization_dyn::visit(reparameterization, visitor);
145                visitor.exit_module(reparameterization.name(), "Reparameterization");
146            }
147        }
148    }
149
150    fn map<M: ModuleMapper>(mut self, mapper: &mut M) -> Self {
151        match self.reparameterization.take() {
152            None => mapper.map_float(self),
153            Some(reparameterization) => {
154                let base = mapper.map_float(self);
155                mapper.enter_module(reparameterization.name(), "Reparameterization");
156                let reparameterization = reparameterization_dyn::map(reparameterization, mapper);
157                mapper.exit_module(reparameterization.name(), "Reparameterization");
158                base.with_dyn_reparameterization(Some(reparameterization))
159            }
160        }
161    }
162
163    fn to_device(mut self, device: &Device) -> Self {
164        let reparameterization = self.reparameterization.take();
165        let base = self.map(|tensor| tensor.to_device(device));
166        match reparameterization {
167            None => base,
168            Some(reparameterization) => {
169                base.with_dyn_reparameterization(Some(reparameterization.to_device_dyn(device)))
170            }
171        }
172    }
173
174    fn fork(mut self, device: &Device) -> Self {
175        let reparameterization = self.reparameterization.take();
176        let base = self.map(|tensor| {
177            let is_require_grad = tensor.is_require_grad();
178            let mut tensor = tensor.to_device(device).detach();
179
180            if is_require_grad {
181                tensor = tensor.require_grad();
182            }
183
184            tensor
185        });
186        match reparameterization {
187            None => base,
188            Some(reparameterization) => {
189                base.with_dyn_reparameterization(Some(reparameterization.fork_dyn(device)))
190            }
191        }
192    }
193
194    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
195        let device = self.base().device();
196
197        if !devices.contains(&device) {
198            devices.push(device)
199        }
200
201        if let Some(reparameterization) = self.reparameterization_dyn() {
202            devices = reparameterization.collect_devices_dyn(devices);
203        }
204
205        devices
206    }
207}
208
209impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D>> {
210    fn content(&self, content: Content) -> Option<Content> {
211        let id = if content.display_settings.show_param_id() {
212            format!(", id: {}", self.id)
213        } else {
214            "".to_string()
215        };
216        let string = format!(
217            "ParamTensor {{rank: {D}, shape: {:?}, kind: float{id}}}",
218            self.shape().as_slice()
219        );
220        content.add_formatted(&string).optional()
221    }
222}
223impl<const D: usize> ModuleDisplay for Param<Tensor<D>> {}
224
225impl<const D: usize> Module for Param<Tensor<D, Int>> {
226    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
227        visitor.visit_int(self)
228    }
229
230    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
231        mapper.map_int(self)
232    }
233
234    fn to_device(self, device: &Device) -> Self {
235        self.map(|tensor| tensor.to_device(device))
236    }
237
238    fn fork(self, device: &Device) -> Self {
239        self.to_device(device) // Don't support autodiff.
240    }
241
242    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
243        let device = self.val().device();
244
245        if !devices.contains(&device) {
246            devices.push(device)
247        }
248
249        devices
250    }
251}
252
253impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Int>> {
254    fn content(&self, content: Content) -> Option<Content> {
255        let id = if content.display_settings.show_param_id() {
256            format!(", id: {}", self.id)
257        } else {
258            "".to_string()
259        };
260        let string = format!(
261            "ParamTensor {{rank: {D}, shape: {:?}, kind: int{id}}}",
262            self.shape().as_slice()
263        );
264        content.add_formatted(&string).optional()
265    }
266}
267impl<const D: usize> ModuleDisplay for Param<Tensor<D, Int>> {}
268
269impl<const D: usize> Module for Param<Tensor<D, Bool>> {
270    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
271        visitor.visit_bool(self)
272    }
273
274    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
275        mapper.map_bool(self)
276    }
277
278    fn to_device(self, device: &Device) -> Self {
279        self.map(|tensor| tensor.to_device(device))
280    }
281
282    fn fork(self, device: &Device) -> Self {
283        self.to_device(device) // Don't support autodiff.
284    }
285
286    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
287        let device = self.val().device();
288
289        if !devices.contains(&device) {
290            devices.push(device)
291        }
292
293        devices
294    }
295}
296
297impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Bool>> {
298    fn content(&self, content: Content) -> Option<Content> {
299        let id = if content.display_settings.show_param_id() {
300            format!(", id: {}", self.id)
301        } else {
302            "".to_string()
303        };
304
305        let string = format!(
306            "ParamTensor {{rank: {D}, shape: {:?}, kind: bool{id}}}",
307            self.shape().as_slice()
308        );
309        content.add_formatted(&string).optional()
310    }
311}
312
313impl<const D: usize> ModuleDisplay for Param<Tensor<D, Bool>> {}
314
315impl<const D: usize> AutodiffModule for Param<Tensor<D>> {
316    fn valid(&self) -> Self {
317        // Preserve initialized param `require_grad` state, but reset the inner value's.
318        // `val()` folds any reparameterization into the base for inference.
319        let require_grad = self.require_grad;
320        let mut param = Param::initialized(self.id, self.val().inner().set_require_grad(false));
321        param.require_grad = require_grad;
322        param
323    }
324
325    fn from_inner(mut module: Self) -> Self {
326        // Keep the reparameterization structure and its parameters on the autodiff backend.
327        let reparameterization = module.reparameterization.take();
328        // Reinstate the param's `require_grad` state
329        let tensor = Tensor::from_inner(module.val()).set_require_grad(module.require_grad);
330        let base = Param::initialized(module.id, tensor);
331        match reparameterization {
332            None => base,
333            Some(reparameterization) => {
334                base.with_dyn_reparameterization(Some(reparameterization.from_inner_dyn()))
335            }
336        }
337    }
338}
339
340impl<const D: usize> AutodiffModule for Param<Tensor<D, Int>> {
341    fn valid(&self) -> Self {
342        Param::initialized(self.id, self.val().inner())
343    }
344
345    fn from_inner(module: Self) -> Self {
346        Param::initialized(module.id, Tensor::from_inner(module.val()))
347    }
348}
349
350impl<const D: usize> AutodiffModule for Param<Tensor<D, Bool>> {
351    fn valid(&self) -> Self {
352        Param::initialized(self.id, self.val().inner())
353    }
354
355    fn from_inner(module: Self) -> Self {
356        Param::initialized(module.id, Tensor::from_inner(module.val()))
357    }
358}
359
360#[cfg(all(test, feature = "std", feature = "autodiff"))]
361mod tests {
362    use super::*;
363    use crate::{module::Module, test_device};
364
365    #[test]
366    fn test_param_require_grad_stateful() {
367        let device = test_device().autodiff();
368        let tensor = Tensor::<2>::ones([3, 3], &device).require_grad();
369
370        let param = Param::initialized(ParamId::new(), tensor);
371        assert!(param.is_require_grad());
372        assert!(param.require_grad);
373
374        let param = param.valid();
375        assert!(!param.is_require_grad());
376        assert!(param.require_grad); // stateful
377
378        // Without `HasAutodiffModule`, we would need to specify the param type as well, which would be annoying:
379        // let param: Param<Tensor<TestAutodiffBackend, _>> = param.train();
380        let param = param.train();
381        assert!(param.is_require_grad());
382        assert!(param.require_grad); // stateful
383
384        let param = param.no_grad();
385        assert!(!param.is_require_grad());
386        assert!(!param.require_grad); // stateful
387
388        let param = param.valid();
389        assert!(!param.is_require_grad()); // always
390        assert!(!param.require_grad); // stateful
391
392        let param = param.train();
393        assert!(!param.is_require_grad());
394        assert!(!param.require_grad); // stateful
395    }
396}
burn_core/module/param/tensor.rs

burn_core/module/param/
tensor.rs
