Struct Mamba1 

pub struct Mamba1 {
    pub in_proj: Linear,
    pub conv1d: Conv1d,
    pub x_proj: Linear,
    pub dt_proj: Linear,
    pub a_log: Param<Tensor<2>>,
    pub d: Param<Tensor<1>>,
    pub out_proj: Linear,
}

The Mamba-1 selective SSM block.

Fields

in_proj: Linear

Input channel: d_model. Output channel: 2 * d_inner.

conv1d: Conv1d

Input channel: d_inner. Output channel: d_inner.

x_proj: Linear

Input channel: d_inner. Output channel: dt_rank + 2 * state_rank.

dt_proj: Linear

Input channel: dt_rank. Output channel: d_inner.

a_log: Param<Tensor<2>>

Dims: [d_inner, state_rank].

d: Param<Tensor<1>>

Dims: [d_inner].

out_proj: Linear

Input channel: d_inner. Output channel: d_model.

Implementations

impl Mamba1

pub fn step( &self, x: Tensor<2>, cache: Option<Mamba1Cache>, ) -> (Tensor<2>, Mamba1Cache)

Shapes

• Input [batch, d_model]
• Output [batch, d_model]

pub fn ssm_step( &self, u: Tensor<2>, cache: Mamba1Cache, ) -> (Tensor<2>, Mamba1Cache)

Single-token counterpart of Mamba1::ssm: computes the selective-SSM parameters (Δ, A, B, C) for one token and advances the recurrence by one step via Self::selective_scan_step.

Shapes

• Input u [batch, d_inner]
• Output [batch, d_inner]

pub fn selective_scan_step( delta: Tensor<2>, a: Tensor<2>, b: Tensor<2>, c: Tensor<2>, d: Tensor<1>, u: Tensor<2>, cache: Mamba1Cache, ) -> (Tensor<2>, Mamba1Cache)

Selective Scan.

Does selective scan algorithm. See:

• Section 2 State Space Models from the Mamba paper;
• Algorithm 2 in Section 3.2 from the Mamba paper;
• run_SSM(A, B, C, u) from The Annotated S4.

Shapes

• Input delta [batch, d_inner]
• Input a [d_inner, state_rank]
• Input b [batch, state_rank]
• Input c [batch, state_rank]
• Input d [d_inner]
• Input u [batch, d_inner]
• Output [batch, d_inner]

impl Mamba1

pub fn forward( &self, x: Tensor<3>, cache: Option<Mamba1Cache>, ) -> (Tensor<3>, Mamba1Cache)

See also Self::step.

Mirrors crate::mamba2::mamba2::Mamba2::forward: an optional cache supplies the initial convolution window and SSM state (zero-initialised when None), and the updated cache is returned so a sequence can be processed in segments (prefill then decode, or chunked prefill).

Shapes

• Input [batch, sequence, d_model]
• Output [batch, sequence, d_model]

pub fn ssm(&self, u: Tensor<3>, init_ssm: Tensor<3>) -> (Tensor<3>, Tensor<3>)

Computes the selective-SSM parameters (Δ, A, B, C) from the conv output and runs the Self::selective_scan recurrence over the full sequence.

Shapes

• Input u [batch, sequence, d_inner]
• Input init_ssm [batch, d_inner, state_rank]
• Output [batch, sequence, d_inner]
• Output (final state) [batch, d_inner, state_rank]

pub fn selective_scan( delta: Tensor<3>, a: Tensor<2>, b: Tensor<3>, c: Tensor<3>, d: Tensor<1>, u: Tensor<3>, init_ssm: Tensor<3>, ) -> (Tensor<3>, Tensor<3>)

Selective Scan.

Does selective scan algorithm. See:

• Section 2 State Space Models from the Mamba paper;
• Algorithm 2 in Section 3.2 from the Mamba paper;
• run_SSM(A, B, C, u) from The Annotated S4.

Shapes

• Input delta [sequence, batch, d_inner]
• Input a [d_inner, state_rank]
• Input b [batch, sequence, state_rank]
• Input c [sequence, batch, state_rank]
• Input d [d_inner]
• Input u [batch, sequence, d_inner]
• Input init_ssm [batch, d_inner, state_rank]
• Output [batch, sequence, d_inner]
• Output (final state) [batch, d_inner, state_rank]

impl Mamba1

fn cache_config(&self, batch: usize) -> Mamba1CacheConfig

fn make_zero( &self, batch: usize, n_virtual: usize, device: &Device, ) -> Mamba1Caches

Trait Implementations

impl AutodiffModule for Mamba1

fn valid(&self) -> Self

Returns the same module, but on the inner backend without auto-differentiation.

fn from_inner(module: Self) -> Self

Wraps an inner module back into an auto-diff module.

impl Clone for Mamba1

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Mamba1

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

Formats the value using the given formatter.

impl Display for Mamba1

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

Formats the value using the given formatter.

impl MambaBlock for Mamba1

type SsdPath = ()

Mamba-1 has no SSD chunking, so there is no path selector.

type Cache = Mamba1Cache

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

type Caches = Mamba1Caches

The per-network cache collection for this family.

fn block_forward( &self, x: Tensor<3>, cache: Option<Mamba1Cache>, _ssd_path: (), ) -> (Tensor<3>, Mamba1Cache)

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

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

Single-token recurrent step — decoding.

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

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

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

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

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

Self::block_forward, additionally returning the exact pooled moments of the block’s per-token SSM states (StateMoments — the inputs of a state participation ratio), matching what a Self::block_step loop reading the cache would accumulate (for Mamba-3, the physical-frame states). The default implementation panics — Mamba-2 provides the closed form (Mamba2::forward_with_state_moments) and Mamba-3 the serial chunkwise de-rotated one (Mamba3::forward_with_state_moments).

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

Self::block_forward_with_state_moments with the moments left attached to the autodiff graph, for a differentiable loss term over them (e.g. a state-PR penalty). The default implementation panics — Mamba-2 and Mamba-3 provide it (see Mamba2::forward_with_state_moments_grad).

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).

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).

impl Module for Mamba1

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices) -> Devices

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn train(self) -> Self

where Self: AutodiffModule,

Move the module and all of its sub-modules to the autodiff backend.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

fn into_record(self) -> ModuleRecord

where Self: Sized,

Collect this module’s parameters into a ModuleRecord.

fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>

where Self: Sized,

Apply a ModuleRecord to this module, returning the loaded module.

fn load_record(self, record: ModuleRecord) -> Self

where Self: Sized,

Apply a ModuleRecord to this module, consuming and returning it.

impl ModuleDisplay for Mamba1

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl ModuleDisplayDefault for Mamba1

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.