Struct SelectiveSSM 

pub struct SelectiveSSM { /* private fields */ }

Available on crate feature alloc only.

Mamba-style selective state space model.

The selective mechanism computes input-dependent B, C, and Delta at each timestep, enabling the model to dynamically control what information is stored in and retrieved from the hidden state.

Dimensions

• d_in – input/output dimension (number of channels)
• n_state – hidden state dimension per channel (N)
• Total hidden state size: d_in * n_state

Weight Shapes

Weight	Shape	Purpose
w_delta	d_in	Projects input to scalar step size
w_b	N x d_in	Projects input to state-input coupling
w_c	N x d_in	Projects input to state-output coupling
d_skip	d_in	Skip connection weights
log_a	N	Fixed state transition (always negative after exp)

Example

use irithyll_core::ssm::selective::SelectiveSSM;
use irithyll_core::ssm::SSMLayer;

let mut ssm = SelectiveSSM::new(4, 8, 42);
let output = ssm.forward(&[1.0, 2.0, 3.0, 4.0]);
assert_eq!(output.len(), 4);

Implementations

impl SelectiveSSM

pub fn new(d_in: usize, n_state: usize, seed: u64) -> Self

Create a new selective SSM with random weight initialization.

Weights are initialized from a small normal distribution (scale 0.1) using the provided seed for reproducibility. A is initialized via the S4D-Inv (HiPPO-inspired) strategy: A_n = -(0.5 + n/N), which gives a bounded spectrum of decay rates that remain meaningful at all state sizes. Skip connections (D) are initialized to 1.0 to enable input passthrough by default.

Arguments

• d_in – input/output dimension (number of channels)
• n_state – hidden state dimension per channel (N)
• seed – random seed for weight initialization

Example

use irithyll_core::ssm::selective::SelectiveSSM;

let ssm = SelectiveSSM::new(4, 16, 42);

pub fn d_in(&self) -> usize

Get the input/output dimension.

pub fn n_state(&self) -> usize

Get the number of state dimensions per channel.

pub fn reinitialize_channel(&mut self, d: usize, rng: &mut u64)

Surgically reinitialize a single channel, preserving all other channels.

Resets channel d’s hidden state to zero across all state dimensions, reinitializes its weight column in w_b and w_c, its w_delta entry, and its skip connection d_skip to the default (1.0). All other channels are left untouched.

Arguments

• d — channel index to reinitialize (must be < d_in)
• rng — mutable RNG state for generating fresh weights

Panics

Panics if d >= d_in.

Trait Implementations

impl SSMLayer for SelectiveSSM

fn forward(&mut self, input: &[f64]) -> Vec<f64>

Process one input timestep and return the output vector.

fn state(&self) -> &[f64]

Get a reference to the current hidden state.

fn output_dim(&self) -> usize

Output dimension of this SSM layer.

fn reset(&mut self)

Reset hidden state to zeros, as if no data has been seen.