burn_tensor/tensor/api/
autodiff.rs

1use crate::{Tensor, kind::Autodiff};
2
3#[cfg(feature = "autodiff")]
4use crate::ops::BridgeTensor;
5#[cfg(feature = "autodiff")]
6use burn_backend::AutodiffBackend;
7#[cfg(feature = "autodiff")]
8use burn_dispatch::Dispatch;
9
10#[cfg(feature = "autodiff")]
11type AutodiffGradients = <Dispatch as AutodiffBackend>::Gradients;
12
13// Aligned, type-erased storage for `AutodiffGradients`. See `crate::macros`
14// for why this indirection exists.
15#[cfg(feature = "autodiff")]
16burn_std::obfuscate!(
17    type: AutodiffGradients,
18    module: gradients_opaque,
19    derives: [Send]
20);
21
22/// Gradients container used during the backward pass.
23#[cfg(feature = "autodiff")]
24pub struct Gradients {
25    blob: gradients_opaque::Opaque,
26}
27
28#[cfg(feature = "autodiff")]
29impl Gradients {
30    /// Crate-internal constructor wrapping the dispatch-level gradients.
31    pub(crate) fn from_inner(inner: AutodiffGradients) -> Self {
32        Self {
33            blob: gradients_opaque::Opaque::new(inner),
34        }
35    }
36
37    /// Crate-internal borrow of the underlying gradients container.
38    pub(crate) fn as_inner(&self) -> &AutodiffGradients {
39        self.blob.as_ref()
40    }
41
42    /// Crate-internal mutable borrow of the underlying gradients container.
43    pub(crate) fn as_inner_mut(&mut self) -> &mut AutodiffGradients {
44        self.blob.as_mut()
45    }
46}
47
48#[cfg(feature = "autodiff")]
49impl<const D: usize> Tensor<D> {
50    /// Backward pass of the tensor.
51    pub fn backward(&self) -> Gradients {
52        backward_impl(&self.primitive)
53    }
54
55    /// Get the gradients of a tensor if it exist.
56    ///
57    /// Returns a new reference to the same tensor. Therefore the same grad tensor can
58    /// be accessed multiple times. If you only need to get the gradients one time,
59    /// consider using [grad_remove](Tensor::grad_remove) for better performance.
60    pub fn grad(&self, grads: &Gradients) -> Option<Tensor<D>> {
61        grad_impl(&self.primitive, grads).map(Tensor::new)
62    }
63
64    /// Remove the grad tensor from the [grads](AutodiffBackend::Gradients) struct returning the result.
65    pub fn grad_remove(&self, grads: &mut Gradients) -> Option<Tensor<D>> {
66        grad_remove_impl(&self.primitive, grads).map(Tensor::new)
67    }
68
69    /// Replace the grad tensor from the [grads](AutodiffBackend::Gradients) struct with the provided
70    /// gradient.
71    pub fn grad_replace(&self, grads: &mut Gradients, grad: Tensor<D>) {
72        grad_replace_impl(&self.primitive, grads, grad.primitive)
73    }
74}
75
76#[cfg(feature = "autodiff")]
77fn backward_impl(p: &BridgeTensor) -> Gradients {
78    Gradients::from_inner(Dispatch::backward(p.clone().into_float()))
79}
80
81#[cfg(feature = "autodiff")]
82fn grad_impl(p: &BridgeTensor, grads: &Gradients) -> Option<BridgeTensor> {
83    Dispatch::grad(p.as_float(), grads.as_inner()).map(BridgeTensor::float)
84}
85
86#[cfg(feature = "autodiff")]
87fn grad_remove_impl(p: &BridgeTensor, grads: &mut Gradients) -> Option<BridgeTensor> {
88    Dispatch::grad_remove(p.as_float(), grads.as_inner_mut()).map(BridgeTensor::float)
89}
90
91#[cfg(feature = "autodiff")]
92fn grad_replace_impl(p: &BridgeTensor, grads: &mut Gradients, grad: BridgeTensor) {
93    Dispatch::grad_replace(p.as_float(), grads.as_inner_mut(), grad.into_float())
94}
95
96impl<const D: usize, K: Autodiff> Tensor<D, K> {
97    /// Returns the inner tensor without the autodiff information.
98    pub fn inner(self) -> Tensor<D, K> {
99        Tensor::new(K::inner(self.primitive))
100    }
101
102    /// Convert a tensor to the autodiff backend.
103    ///
104    /// # Arguments
105    ///
106    /// * `inner` - The tensor to convert.
107    ///
108    /// # Returns
109    ///
110    /// The tensor converted to the autodiff backend.
111    pub fn from_inner(inner: Tensor<D, K>) -> Self {
112        Self::new(K::from_inner(inner.primitive))
113    }
114}
115
116// TODO: a lot of the `tensor.inner` / `Tensor::from_inner(...)` are actually scoped to perform some operations
117// so it might be cleaner and easier to manage the device etc. if we provide a method to scope the autodiff?
burn_tensor/tensor/api/autodiff.rs

burn_tensor/tensor/api/
autodiff.rs
