Struct CheckpointBudget 

pub struct CheckpointBudget {
    pub max_bytes: usize,
    pub used_bytes: usize,
}

Memory budget tracker for gradient checkpointing.

Maintains a running total of bytes allocated to checkpointed inputs and enforces an upper bound. Call allocate when a new checkpoint is created and free when it is discarded (e.g., after the backward pass for that segment completes).

Fields

max_bytes: usize

Maximum permitted allocation in bytes.

used_bytes: usize

Currently allocated bytes.

Implementations

impl CheckpointBudget

pub fn new(max_bytes: usize) -> Self

Create a fresh budget with max_bytes capacity and zero usage.

pub fn remaining(&self) -> usize

Bytes still available for allocation.

pub fn utilization(&self) -> f32

Fraction of the budget that has been consumed (used / max).

Returns 0.0 when max_bytes == 0 to avoid division by zero.

pub fn can_allocate(&self, bytes: usize) -> bool

Whether bytes can be allocated without exceeding the budget.

pub fn allocate(&mut self, bytes: usize) -> Result<(), CheckpointError>

Attempt to allocate bytes.

On success, used_bytes increases by bytes. On failure, returns CheckpointError::BudgetExceeded and leaves used_bytes unchanged.

pub fn free(&mut self, bytes: usize)

Release bytes back to the budget.

Uses saturating subtraction to avoid underflow if bytes exceeds used_bytes (which would indicate a programming error, but should not panic in production).

Trait Implementations

impl Clone for CheckpointBudget

fn clone(&self) -> CheckpointBudget

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for CheckpointBudget

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

Formats the value using the given formatter.