Struct SelectiveSSMBD 

pub struct SelectiveSSMBD { /* private fields */ }

Available on crate feature alloc only.

Block-Diagonal Linear Recurrent Unit selective state space model.

Partitions d_in channels into n_blocks = d_in / block_size blocks, each with a dense block_size x block_size A matrix for within-block cross-channel state mixing. B, C, and Delta projections are shared across blocks (same structure as Mamba-1).

Dimensions

• d_in – input/output dimension (number of channels)
• n_state – hidden state dimension per block-channel (N)
• block_size – number of channels per block (m)
• n_blocks – number of blocks (d_in / block_size)
• Total hidden state size: n_blocks * n_state * block_size

Weight Shapes

Weight	Shape	Purpose
a_matrices	n_blocks * m * m	Dense A per block (row-major, L1-normalized)
w_b	N x d_in	Projects input to state-input coupling
w_c	N x d_in	Projects input to state-output coupling
w_delta	d_in	Projects input to scalar step size
d_skip	d_in	Skip connection weights

Example

use irithyll_core::ssm::selective_bd::SelectiveSSMBD;
use irithyll_core::ssm::SSMLayer;

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

Implementations

impl SelectiveSSMBD

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

Create a new block-diagonal selective SSM with random weight initialization.

A matrices are initialized with S4D-Inv diagonal values and small random off-diagonal entries (scale 0.02), then row-wise L1-normalized. Projection weights are initialized from a small normal distribution (scale 0.1). Skip connections (D) are initialized to 1.0 for input passthrough.

Arguments

• d_in – input/output dimension (must be divisible by block_size)
• n_state – hidden state dimension per block-channel (N)
• block_size – number of channels per block (m)
• seed – random seed for weight initialization

Panics

Panics if d_in is not evenly divisible by block_size.

Example

use irithyll_core::ssm::selective_bd::SelectiveSSMBD;

let ssm = SelectiveSSMBD::new(6, 8, 2, 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 block-channel.

pub fn block_size(&self) -> usize

Get the number of channels per block.

pub fn n_blocks(&self) -> usize

Get the number of blocks.

pub fn reinitialize_block(&mut self, b: usize, rng: &mut u64)

Surgically reinitialize a single block, preserving all other blocks.

Resets block b’s hidden state to zero, reinitializes its A matrix with S4D diagonal + small random off-diagonal values (then L1 row- normalizes), and resets the skip connections for the block’s channels to 1.0. All other blocks are left untouched.

Arguments

• b — block index to reinitialize (must be < n_blocks)
• rng — mutable RNG state for generating fresh weights

Panics

Panics if b >= n_blocks.

Trait Implementations

impl SSMLayer for SelectiveSSMBD

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.