Skip to main content

CacheStack

burn_mamba::prelude

Trait CacheStack 

pub trait CacheStack: Sized {
    type Cache;

    // Required methods
    fn slot_count(&self) -> usize;
    fn into_slots(self) -> Vec<Option<Self::Cache>>;
    fn from_slots(slots: Vec<Option<Self::Cache>>) -> Self;
    fn cache_to_inner(cache: Self::Cache) -> Self::Cache;
    fn cache_from_inner(cache: Self::Cache) -> Self::Cache;
}
Expand description

The uniform interface a per-network cache collection exposes for the generic Layers loop: slot_count + move-in/move-out of the per-layer slots.

Required Associated Types§

type Cache

The per-layer cache element.

Required Methods§

fn slot_count(&self) -> usize

Number of per-(virtual-)layer slots.

fn into_slots(self) -> Vec<Option<Self::Cache>>

Move each slot into an Option so the loop can take without cloning.

fn from_slots(slots: Vec<Option<Self::Cache>>) -> Self

Inverse of Self::into_slots.

fn cache_to_inner(cache: Self::Cache) -> Self::Cache

Move one cache slot to the inner (non-autodiff) backend.

Needed by Layers::grad_horizon, whose no-grad prefix runs on the inner backend: a cache carried in from a tracked segment has to come down with it, both so the prefix builds no graph and because Burn’s dispatch cannot mix backends within one op.

Spelled out per family rather than derived, because Module::map is a no-op on plain Tensor fields (Burn implements Module for Tensor as a constant) and caches hold bare tensors, not Params — a Module-based conversion would silently skip every one of them.

§Panics

Tensor::inner panics on a tensor that is already off the autodiff backend, so the caller must have checked Device::is_autodiff first.

fn cache_from_inner(cache: Self::Cache) -> Self::Cache

Lift one cache slot back from the inner backend, as a fresh graph root. The inverse of Self::cache_to_inner; see its notes.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§