Struct burn::backend::autodiff::grads::Gradients

pub struct Gradients { /* private fields */ }

Gradients container used during the backward pass.

Implementations

impl Gradients

pub fn new<B>( root_node: Arc<Node>, root_tensor: <B as Backend>::FloatTensorPrimitive, ) -> Gradients

where B: Backend,

Creates a new gradients container.

pub fn consume<B>( &mut self, node: &Arc<Node>, ) -> <B as Backend>::FloatTensorPrimitive

where B: Backend,

Consumes the gradients for a given tensor.

Each tensor should be consumed exactly 1 time if its gradients are only required during the backward pass, otherwise, it may be consume multiple times.

pub fn remove<B>( &mut self, tensor: &AutodiffTensor<B>, ) -> Option<<B as Backend>::FloatTensorPrimitive>

where B: Backend,

Removes a grad tensor from the container.

pub fn get<B>( &self, tensor: &AutodiffTensor<B>, ) -> Option<<B as Backend>::FloatTensorPrimitive>

where B: Backend,

Gets a grad tensor from the container.

pub fn register<B>( &mut self, node_id: NodeID, value: <B as Backend>::FloatTensorPrimitive, )

where B: Backend,

Register a grad tensor in the container.

If the tensor already exists, add both tensors together before saving the result.