Skip to main content

Block

burn_mamba::prelude

Trait Block 

pub trait Block:
    Module
    + ModuleDisplay
    + AutodiffModule {
    type Cache;
    type Caches: CacheStack<Cache = Self::Cache>;
    type Options;

    // Required methods
    fn block_forward(
        &self,
        x: Tensor<3>,
        cache: Option<Self::Cache>,
        options: Self::Options,
    ) -> (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 method
    fn block_step_infinite(&self, x: Tensor<2>) -> Tensor<2> { ... }
}
Expand description

The mixer-block interface the generic Layer/Layers delegate to.

Implement it once per block family (a selective SSM, an attention variant, a gated convolution, …) and every container in this crate — layers, virtual stacks, bidirectional pairs, latent/vocab networks, class tokens, the Muon plan — applies unchanged.

ModuleDisplay and AutodiffModule are supertraits so that the generic containers are themselves Module/AutodiffModule (Burn’s derive requires both of every module-typed generic), which is what lets Layers::grad_horizon move the stack to the inner backend for its no-grad prefix. A #[derive(Module)] block satisfies them.

Required Associated Types§

type Cache

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

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

The per-network cache collection for this family.

type Options

Per-call algorithm/chunking options threaded down to Self::block_forward. () for a block with nothing to select.

Required Methods§

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

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

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

Single-token recurrent step — decoding.

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.

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§

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 — a block only provides this when its recurrence has a closed-form constant-input limit.

Dyn Compatibility§

This trait is not dyn compatible.

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

Implementors§