Struct 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 BackendTypes>::FloatTensorPrimitive, ) -> Gradients

where B: Backend,

Creates a new gradients container.

pub fn new_distributed<B>( root_node: Arc<Node>, root_tensor: <B as BackendTypes>::FloatTensorPrimitive, reg: Box<dyn DistributedRegistration>, ) -> Gradients

where B: Backend,

Creates a new gradients container with a registration hook for distributed gradients.

pub fn consume<B>( &mut self, node: &Arc<Node>, ) -> <B as BackendTypes>::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 BackendTypes>::FloatTensorPrimitive>

where B: Backend,

Removes a grad tensor from the container.

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

where B: Backend,

Gets a grad tensor from the container.

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

where B: Backend,

Register a grad tensor in the container.

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

If the registered tensor is distributed, launches a syncing operation on the gradients.