Struct ComputationGraph 

pub struct ComputationGraph { /* private fields */ }

Computation graph that records operations for backward pass.

The graph uses a tape-based approach where operations are recorded in order during the forward pass, then gradients are computed in reverse order during the backward pass.

Thread Safety

Each thread has its own computation graph (via thread_local storage in the parent module). This avoids synchronization overhead during single-threaded training.

Implementations

impl ComputationGraph

pub fn new() -> ComputationGraph

Create a new empty computation graph.

pub fn clear(&mut self)

Clear all recorded operations.

pub fn register_tensor(&mut self, tensor: Tensor)

Register a tensor that requires gradients.

pub fn record( &mut self, output_id: TensorId, grad_fn: Arc<dyn GradFn>, input_ids: Vec<TensorId>, )

Record an operation to the tape.

pub fn get_tensor(&self, id: TensorId) -> Option<&Tensor>

Get a tensor by ID.

pub fn get_tensor_mut(&mut self, id: TensorId) -> Option<&mut Tensor>

Get a mutable tensor by ID.

pub fn backward(&mut self, output_id: TensorId, grad_output: Tensor)

Compute gradients via backpropagation.

This implements the reverse-mode automatic differentiation algorithm:

1. Start with grad_output for the output tensor
2. Iterate through operations in reverse order
3. For each operation, compute gradients w.r.t. inputs
4. Accumulate gradients for tensors used multiple times

Arguments

• output_id - ID of the tensor to differentiate
• grad_output - Initial gradient (typically ones for scalar loss)

pub fn len(&self) -> usize

Get the number of recorded operations.

pub fn is_empty(&self) -> bool

Check if the tape is empty.

pub fn get_grad(&self, id: TensorId) -> Option<Tensor>

Get gradient for a tensor by ID (after backward).

pub fn clear_grad(&mut self, id: TensorId)

Clear gradient for a specific tensor.

Trait Implementations

impl Default for ComputationGraph

fn default() -> ComputationGraph

Returns the “default value” for a type.