Struct ActivationCheckpointer 

pub struct ActivationCheckpointer { /* private fields */ }

Manages activation checkpoints for a multi-layer forward pass.

Stores activations at designated layer boundaries so they can be used during backpropagation without keeping every intermediate result in memory.

Implementations

impl ActivationCheckpointer

pub fn new(config: CheckpointConfig) -> Self

Create a new checkpointer with the given configuration.

pub fn save( &mut self, layer_idx: usize, activation: Array1<f32>, ) -> ModelResult<()>

Save an activation at the given layer index.

If use_mixed_precision is enabled, the activation is quantised before storage.

Errors

Returns an error if max_checkpoints would be exceeded and the layer is not a checkpoint boundary.

pub fn get(&self, layer_idx: usize) -> ModelResult<&Array1<f32>>

Retrieve the checkpointed activation at the given layer.

Errors

Returns an error if no checkpoint exists for layer_idx.

pub fn clear(&mut self)

Clear all stored checkpoints and reset memory accounting.

pub fn memory_saved_bytes(&self) -> usize

Estimated bytes of memory saved by not storing non-checkpointed activations.

This value is updated during checkpointed_forward calls.

pub fn memory_stored_bytes(&self) -> usize

Bytes currently stored in checkpoints.

pub fn num_checkpoints(&self) -> usize

Number of non-None checkpoints currently held.

pub fn is_checkpoint_layer(&self, layer_idx: usize) -> bool

Whether a given layer index is a checkpoint boundary according to the current configuration.

pub fn checkpointed_forward<F>( &mut self, input: &Array1<f32>, layers: &[usize], forward_fn: F, ) -> ModelResult<Array1<f32>>

where F: Fn(&Array1<f32>, usize) -> ModelResult<Array1<f32>>,

Run a checkpointed forward pass through the given layers.

The forward_fn is called sequentially for each layer in layers, receiving the current activation and the layer index. Activations at checkpoint boundaries are saved; others are discarded (their memory cost is recorded in bytes_saved).

Parameters

• input: the initial activation fed into the first layer.
• layers: ordered list of layer indices to process.
• forward_fn: Fn(&Array1<f32>, usize) -> ModelResult<Array1<f32>>; applies one layer’s computation.

Returns

The activation after all layers have been applied.

pub fn recompute_from_checkpoint<F>( &self, target_layer: usize, layers: &[usize], forward_fn: F, ) -> ModelResult<Array1<f32>>

where F: Fn(&Array1<f32>, usize) -> ModelResult<Array1<f32>>,

Recompute activations from the nearest checkpoint up to target_layer.

This is used during the backward pass: find the closest checkpoint before target_layer, then replay the forward function from there.

Parameters

• target_layer: the layer whose activation is needed.
• layers: the full ordered list of layer indices.
• forward_fn: the same forward function used during the forward pass.

Returns

The recomputed activation at target_layer.

pub fn config(&self) -> &CheckpointConfig

Return the configuration.

Trait Implementations

impl Clone for ActivationCheckpointer

fn clone(&self) -> ActivationCheckpointer

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ActivationCheckpointer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.