Skip to main content

MambaBlock

burn_mamba::modules

Trait MambaBlock 

Source 
pub trait MambaBlock: Module {
    type Cache;
    type Caches: CacheStack<Cache = Self::Cache>;
    type SsdPath;

    // Required methods
    fn block_forward(
        &self,
        x: Tensor<3>,
        cache: Option<Self::Cache>,
        ssd_path: Self::SsdPath,
    ) -> (Tensor<3>, Self::Cache);
    fn block_step(
        &self,
        x: Tensor<2>,
        cache: Option<Self::Cache>,
    ) -> (Tensor<2>, Self::Cache);
    fn zero_caches_3d(&self, x: &Tensor<3>, n_virtual: usize) -> Self::Caches;
    fn zero_caches_2d(&self, x: &Tensor<2>, n_virtual: usize) -> Self::Caches;

    // Provided methods
    fn block_step_infinite(&self, x: Tensor<2>) -> Tensor<2> { ... }
    fn block_step_n_approx(
        &self,
        x: Tensor<2>,
        n: usize,
        cache: Option<Self::Cache>,
    ) -> (Tensor<2>, Self::Cache) { ... }
}
Expand description

Per-family block interface the generic Layer/Layers delegate to.

Required Associated Types§

Source

type Cache

Per-block streaming cache (one layer’s worth of state).

Source

type Caches: CacheStack<Cache = Self::Cache>

The per-network cache collection for this family.

Source

type SsdPath

SSD algorithm / chunk-length selector. () for families without one.

Required Methods§

Source

fn block_forward( &self, x: Tensor<3>, cache: Option<Self::Cache>, ssd_path: Self::SsdPath, ) -> (Tensor<3>, Self::Cache)

Full-sequence (chunked) pass — training / prefill.

Source

fn block_step( &self, x: Tensor<2>, cache: Option<Self::Cache>, ) -> (Tensor<2>, Self::Cache)

Single-token recurrent step — decoding.

Source

fn zero_caches_3d(&self, x: &Tensor<3>, n_virtual: usize) -> Self::Caches

Build n_virtual zero caches sized for a [batch, sequence, d_model] input.

Source

fn zero_caches_2d(&self, x: &Tensor<2>, n_virtual: usize) -> Self::Caches

Build n_virtual zero caches sized for a [batch, d_model] input.

Provided Methods§

Source

fn block_step_infinite(&self, x: Tensor<2>) -> Tensor<2>

Closed-form stationary fixed point: the limit of Self::block_step outputs when the same constant token is stepped forever. The limit forgets the starting state, so no cache is taken or returned. The default implementation panics — only Mamba-3 currently provides the closed form (see Mamba3::step_infinite).

Source

fn block_step_n_approx( &self, x: Tensor<2>, n: usize, cache: Option<Self::Cache>, ) -> (Tensor<2>, Self::Cache)

Closed-form jump equivalent to n consecutive Self::block_step calls on the same constant token: the last step’s output and the final cache, in O(1). The default implementation panics — only Mamba-3 currently provides it (see Mamba3::step_n_approx).

Dyn Compatibility§

This trait is not dyn compatible.

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

Implementors§