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

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

//! Inspect the tensor keys / dtypes / shapes of a safetensors checkpoint.
//!
//! Usage: `safetensors_info <path.safetensors>`

use std::collections::BTreeMap;

use safetensors::SafeTensors;

fn main() -> anyhow::Result<()> {
    let path = std::env::args()
        .nth(1)
        .ok_or_else(|| anyhow::anyhow!("usage: safetensors_info <path.safetensors>"))?;
    let bytes = std::fs::read(&path)?;
    let st = SafeTensors::deserialize(&bytes)?;

    let mut entries: BTreeMap<String, (String, Vec<usize>, usize)> = BTreeMap::new();
    let mut total: usize = 0;
    for (k, v) in st.tensors() {
        let n: usize = v.shape().iter().product::<usize>().max(1);
        total += n;
        entries.insert(k.clone(), (format!("{:?}", v.dtype()), v.shape().to_vec(), n));
    }

    println!("{}  ({} tensors)", path, entries.len());
    for (k, (dt, shape, _)) in &entries {
        println!("  {k:60}  {dt:>5}  {shape:?}");
    }
    println!("total parameters: {total} ({:.2} M)", total as f64 / 1e6);
    Ok(())
}