lumamba 0.0.1

LuMamba EEG foundation model — inference in Rust on the RLX runtime
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

// lumamba-rs — LuMamba EEG foundation-model inference on the RLX runtime.
// Copyright (C) 2026 Nataliya Kosmyna.
// SPDX-License-Identifier: GPL-3.0-only

//! # lumamba-rs — LuMamba EEG foundation-model inference in Rust
//!
//! Pure-Rust inference for **LuMamba** ([`PulpBio/LuMamba`](https://huggingface.co/PulpBio/LuMamba),
//! from [BioFoundation](https://github.com/pulp-bio/BioFoundation)), built on the
//! [RLX](https://github.com/MIT-RLX/rlx) compiler/runtime.
//!
//! LuMamba reuses LUNA's topology-invariant front-end — variable-channel EEG is
//! compressed into a fixed set of learned queries by cross-attention — and replaces
//! LUNA's rotary Transformer encoder with a stack of **bidirectional Mamba (FEMBA)
//! blocks** that run in linear time over the patch sequence.
//!
//! ```text
//! EEG (B, C, T)
//!   ├─ PatchEmbed (3× Conv2d + GroupNorm + GELU)         ┐  CPU prepare
//!   └─ FreqEmbed  (FFT magnitude/phase → MLP)            ├─ (host)
//!              → + NeRF channel-location embedding       ┘
//!   → CrossAttention channel unification (C → Q queries)  ┐
//!   → reshape [B, S, Q·E]                                 │  RLX graph
//!   → N × { LayerNorm → BiMamba(fwd + flip∘rev) → +res }  │
//!   → reconstruction head  →  (B, C, T)                   ┘
//! ```
//!
//! The temporal recurrence runs on RLX's first-class [`selective_scan`] op
//! (Mamba-1 SSM), the depthwise causal conv on grouped `conv2d`, and everything
//! else on the standard graph ops — so the whole model compiles to a single
//! RLX `CompiledGraph` per input shape and runs on any RLX backend
//! (`cpu`, `metal`, `mlx`, `cuda`, …).
//!
//! [`selective_scan`]: https://github.com/MIT-RLX/rlx
#![warn(missing_docs)]
// Indexing by range index is often clearer in the numeric / tensor-shape code
// here (matches the RLX workspace's clippy posture).
#![allow(clippy::needless_range_loop)]

/// Configure the global Rayon thread pool (used by the CPU prepare path and
/// the RLX CPU backend).
pub fn init_threads(n: Option<usize>) -> usize {
    let mut builder = rayon::ThreadPoolBuilder::new();
    if let Some(count) = n {
        if count > 0 {
            builder = builder.num_threads(count);
        }
    }
    let _ = builder.build_global();
    rayon::current_num_threads()
}

/// The RLX backends lumamba-rs forwards Cargo features for. A `--device`
/// token outside this set is rejected at parse time by [`parse_device`].
///
/// (`::rlx` is the external runtime crate — this crate's own `rlx` module
/// shadows the bare name.)
pub const SUPPORTED_DEVICES: &[::rlx::Device] = &[
    ::rlx::Device::Cpu,
    ::rlx::Device::Metal,
    ::rlx::Device::Mlx,
    ::rlx::Device::Gpu,
    ::rlx::Device::Cuda,
    ::rlx::Device::Rocm,
    ::rlx::Device::Tpu,
];

/// [`BackendSupport`](::rlx::driver::BackendSupport) for the LuMamba model
/// family — the RLX 0.2.10 idiom for declaring which devices a model can
/// execute on, so device validation yields a uniform error instead of a
/// hand-rolled `match` ladder.
#[derive(Debug, Clone, Copy)]
pub struct LuMambaBackends;

impl ::rlx::driver::BackendSupport for LuMambaBackends {
    fn family(&self) -> &'static str {
        "lumamba"
    }
    fn supports(&self, device: ::rlx::Device) -> bool {
        SUPPORTED_DEVICES.contains(&device)
    }
}

/// Parse a `--device` CLI token into an [`rlx::Device`](::rlx::Device).
///
/// Uses RLX 0.2.10's [`FromStr`](std::str::FromStr) (so aliases like `mps` /
/// `wgpu` resolve), then gates to [`SUPPORTED_DEVICES`] via
/// [`validate_device`](::rlx::driver::validate_device). Designed as a clap
/// `value_parser`.
pub fn parse_device(s: &str) -> Result<::rlx::Device, String> {
    let device = s.parse::<::rlx::Device>().map_err(|e| e.to_string())?;
    ::rlx::driver::validate_device(&LuMambaBackends, device)
}

pub mod channel_positions;
pub mod channel_vocab;
pub mod config;
pub mod eval;
pub mod hf;
pub mod metrics;
pub mod rlx;

pub use config::ModelConfig;

pub use rlx::{ClassifierKind, EpochEmbedding, LuMambaEncoder, RlxEpoch, RunEpochOpts};

pub use channel_positions::{
    bipolar_channel_xyz, channel_xyz, montage_channels, nearest_channel, normalise,
    MontageLayout,
};

pub use channel_vocab::{
    channel_index, channel_indices, channel_indices_unwrap, CHANNEL_VOCAB, SEED_CHANNELS,
    SIENA_CHANNELS, TUEG_CHANNELS, VOCAB_SIZE,
};