383 lines
12 KiB
Rust
383 lines
12 KiB
Rust
//! Unified WebAssembly Attention Library
|
|
//!
|
|
//! This crate provides a unified WASM interface for 18+ attention mechanisms:
|
|
//!
|
|
//! ## Neural Attention (from ruvector-attention)
|
|
//! - **Scaled Dot-Product**: Standard transformer attention
|
|
//! - **Multi-Head**: Parallel attention heads
|
|
//! - **Hyperbolic**: Attention in hyperbolic space for hierarchical data
|
|
//! - **Linear**: O(n) Performer-style attention
|
|
//! - **Flash**: Memory-efficient blocked attention
|
|
//! - **Local-Global**: Sparse attention with global tokens
|
|
//! - **MoE**: Mixture of Experts attention
|
|
//!
|
|
//! ## DAG Attention (from ruvector-dag)
|
|
//! - **Topological**: Position-aware attention in DAG order
|
|
//! - **Causal Cone**: Lightcone-based causal attention
|
|
//! - **Critical Path**: Attention weighted by critical path distance
|
|
//! - **MinCut-Gated**: Flow-based gating attention
|
|
//! - **Hierarchical Lorentz**: Multi-scale hyperbolic DAG attention
|
|
//! - **Parallel Branch**: Attention for parallel DAG branches
|
|
//! - **Temporal BTSP**: Behavioral Time-Series Pattern attention
|
|
//!
|
|
//! ## Graph Attention (from ruvector-gnn)
|
|
//! - **GAT**: Graph Attention Networks
|
|
//! - **GCN**: Graph Convolutional Networks
|
|
//! - **GraphSAGE**: Sampling and Aggregating graph embeddings
|
|
//!
|
|
//! ## State Space Models
|
|
//! - **Mamba SSM**: Selective State Space Model attention
|
|
|
|
use wasm_bindgen::prelude::*;
|
|
|
|
// Use wee_alloc for smaller WASM binary (~10KB reduction)
|
|
#[cfg(feature = "wee_alloc")]
|
|
#[global_allocator]
|
|
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
|
|
|
|
// ============================================================================
|
|
// Module declarations
|
|
// ============================================================================
|
|
|
|
pub mod mamba;
|
|
|
|
mod dag;
|
|
mod graph;
|
|
mod neural;
|
|
|
|
// ============================================================================
|
|
// Re-exports for convenient access
|
|
// ============================================================================
|
|
|
|
pub use dag::*;
|
|
pub use graph::*;
|
|
pub use mamba::*;
|
|
pub use neural::*;
|
|
|
|
// ============================================================================
|
|
// Initialization
|
|
// ============================================================================
|
|
|
|
/// Initialize the WASM module with panic hook for better error messages
|
|
#[wasm_bindgen(start)]
|
|
pub fn init() {
|
|
#[cfg(feature = "console_error_panic_hook")]
|
|
console_error_panic_hook::set_once();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Version and Info
|
|
// ============================================================================
|
|
|
|
/// Get the version of the unified attention WASM crate
|
|
#[wasm_bindgen]
|
|
pub fn version() -> String {
|
|
env!("CARGO_PKG_VERSION").to_string()
|
|
}
|
|
|
|
/// Get information about all available attention mechanisms
|
|
#[wasm_bindgen(js_name = availableMechanisms)]
|
|
pub fn available_mechanisms() -> JsValue {
|
|
let mechanisms = AttentionMechanisms {
|
|
neural: vec![
|
|
"scaled_dot_product".into(),
|
|
"multi_head".into(),
|
|
"hyperbolic".into(),
|
|
"linear".into(),
|
|
"flash".into(),
|
|
"local_global".into(),
|
|
"moe".into(),
|
|
],
|
|
dag: vec![
|
|
"topological".into(),
|
|
"causal_cone".into(),
|
|
"critical_path".into(),
|
|
"mincut_gated".into(),
|
|
"hierarchical_lorentz".into(),
|
|
"parallel_branch".into(),
|
|
"temporal_btsp".into(),
|
|
],
|
|
graph: vec!["gat".into(), "gcn".into(), "graphsage".into()],
|
|
ssm: vec!["mamba".into()],
|
|
};
|
|
serde_wasm_bindgen::to_value(&mechanisms).unwrap()
|
|
}
|
|
|
|
/// Get summary statistics about the unified attention library
|
|
#[wasm_bindgen(js_name = getStats)]
|
|
pub fn get_stats() -> JsValue {
|
|
let stats = UnifiedStats {
|
|
total_mechanisms: 18,
|
|
neural_count: 7,
|
|
dag_count: 7,
|
|
graph_count: 3,
|
|
ssm_count: 1,
|
|
version: env!("CARGO_PKG_VERSION").to_string(),
|
|
};
|
|
serde_wasm_bindgen::to_value(&stats).unwrap()
|
|
}
|
|
|
|
// ============================================================================
|
|
// Internal Types
|
|
// ============================================================================
|
|
|
|
#[derive(serde::Serialize)]
|
|
struct AttentionMechanisms {
|
|
neural: Vec<String>,
|
|
dag: Vec<String>,
|
|
graph: Vec<String>,
|
|
ssm: Vec<String>,
|
|
}
|
|
|
|
#[derive(serde::Serialize)]
|
|
struct UnifiedStats {
|
|
total_mechanisms: usize,
|
|
neural_count: usize,
|
|
dag_count: usize,
|
|
graph_count: usize,
|
|
ssm_count: usize,
|
|
version: String,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Unified Attention Selector
|
|
// ============================================================================
|
|
|
|
/// Unified attention mechanism selector
|
|
/// Automatically routes to the appropriate attention implementation
|
|
#[wasm_bindgen]
|
|
pub struct UnifiedAttention {
|
|
mechanism_type: String,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl UnifiedAttention {
|
|
/// Create a new unified attention selector
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn new(mechanism: &str) -> Result<UnifiedAttention, JsError> {
|
|
let valid_mechanisms = [
|
|
// Neural
|
|
"scaled_dot_product",
|
|
"multi_head",
|
|
"hyperbolic",
|
|
"linear",
|
|
"flash",
|
|
"local_global",
|
|
"moe",
|
|
// DAG
|
|
"topological",
|
|
"causal_cone",
|
|
"critical_path",
|
|
"mincut_gated",
|
|
"hierarchical_lorentz",
|
|
"parallel_branch",
|
|
"temporal_btsp",
|
|
// Graph
|
|
"gat",
|
|
"gcn",
|
|
"graphsage",
|
|
// SSM
|
|
"mamba",
|
|
];
|
|
|
|
if !valid_mechanisms.contains(&mechanism) {
|
|
return Err(JsError::new(&format!(
|
|
"Unknown mechanism: {}. Valid options: {:?}",
|
|
mechanism, valid_mechanisms
|
|
)));
|
|
}
|
|
|
|
Ok(Self {
|
|
mechanism_type: mechanism.to_string(),
|
|
})
|
|
}
|
|
|
|
/// Get the currently selected mechanism type
|
|
#[wasm_bindgen(getter)]
|
|
pub fn mechanism(&self) -> String {
|
|
self.mechanism_type.clone()
|
|
}
|
|
|
|
/// Get the category of the selected mechanism
|
|
#[wasm_bindgen(getter)]
|
|
pub fn category(&self) -> String {
|
|
match self.mechanism_type.as_str() {
|
|
"scaled_dot_product" | "multi_head" | "hyperbolic" | "linear" | "flash"
|
|
| "local_global" | "moe" => "neural".to_string(),
|
|
|
|
"topological"
|
|
| "causal_cone"
|
|
| "critical_path"
|
|
| "mincut_gated"
|
|
| "hierarchical_lorentz"
|
|
| "parallel_branch"
|
|
| "temporal_btsp" => "dag".to_string(),
|
|
|
|
"gat" | "gcn" | "graphsage" => "graph".to_string(),
|
|
|
|
"mamba" => "ssm".to_string(),
|
|
|
|
_ => "unknown".to_string(),
|
|
}
|
|
}
|
|
|
|
/// Check if this mechanism supports sequence processing
|
|
#[wasm_bindgen(js_name = supportsSequences)]
|
|
pub fn supports_sequences(&self) -> bool {
|
|
matches!(
|
|
self.mechanism_type.as_str(),
|
|
"scaled_dot_product" | "multi_head" | "linear" | "flash" | "local_global" | "mamba"
|
|
)
|
|
}
|
|
|
|
/// Check if this mechanism supports graph/DAG structures
|
|
#[wasm_bindgen(js_name = supportsGraphs)]
|
|
pub fn supports_graphs(&self) -> bool {
|
|
matches!(
|
|
self.mechanism_type.as_str(),
|
|
"topological"
|
|
| "causal_cone"
|
|
| "critical_path"
|
|
| "mincut_gated"
|
|
| "hierarchical_lorentz"
|
|
| "parallel_branch"
|
|
| "temporal_btsp"
|
|
| "gat"
|
|
| "gcn"
|
|
| "graphsage"
|
|
)
|
|
}
|
|
|
|
/// Check if this mechanism supports hyperbolic geometry
|
|
#[wasm_bindgen(js_name = supportsHyperbolic)]
|
|
pub fn supports_hyperbolic(&self) -> bool {
|
|
matches!(
|
|
self.mechanism_type.as_str(),
|
|
"hyperbolic" | "hierarchical_lorentz"
|
|
)
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Utility Functions
|
|
// ============================================================================
|
|
|
|
/// Compute cosine similarity between two vectors
|
|
#[wasm_bindgen(js_name = cosineSimilarity)]
|
|
pub fn cosine_similarity(a: Vec<f32>, b: Vec<f32>) -> Result<f32, JsError> {
|
|
if a.len() != b.len() {
|
|
return Err(JsError::new(&format!(
|
|
"Vector dimensions must match: {} vs {}",
|
|
a.len(),
|
|
b.len()
|
|
)));
|
|
}
|
|
|
|
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
|
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
|
|
if norm_a == 0.0 || norm_b == 0.0 {
|
|
Ok(0.0)
|
|
} else {
|
|
Ok(dot / (norm_a * norm_b))
|
|
}
|
|
}
|
|
|
|
/// Softmax normalization
|
|
#[wasm_bindgen]
|
|
pub fn softmax(values: Vec<f32>) -> Vec<f32> {
|
|
let max_val = values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
let exp_values: Vec<f32> = values.iter().map(|&x| (x - max_val).exp()).collect();
|
|
let sum: f32 = exp_values.iter().sum();
|
|
exp_values.iter().map(|&x| x / sum).collect()
|
|
}
|
|
|
|
/// Temperature-scaled softmax
|
|
#[wasm_bindgen(js_name = temperatureSoftmax)]
|
|
pub fn temperature_softmax(values: Vec<f32>, temperature: f32) -> Vec<f32> {
|
|
if temperature <= 0.0 {
|
|
// Return one-hot for the maximum
|
|
let max_idx = values
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
|
.map(|(i, _)| i)
|
|
.unwrap_or(0);
|
|
let mut result = vec![0.0; values.len()];
|
|
result[max_idx] = 1.0;
|
|
return result;
|
|
}
|
|
|
|
let scaled: Vec<f32> = values.iter().map(|&x| x / temperature).collect();
|
|
softmax(scaled)
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use wasm_bindgen_test::*;
|
|
|
|
wasm_bindgen_test_configure!(run_in_browser);
|
|
|
|
#[wasm_bindgen_test]
|
|
fn test_version() {
|
|
assert!(!version().is_empty());
|
|
}
|
|
|
|
#[wasm_bindgen_test]
|
|
fn test_unified_attention_creation() {
|
|
let attention = UnifiedAttention::new("multi_head");
|
|
assert!(attention.is_ok());
|
|
|
|
let invalid = UnifiedAttention::new("invalid_mechanism");
|
|
assert!(invalid.is_err());
|
|
}
|
|
|
|
#[wasm_bindgen_test]
|
|
fn test_mechanism_categories() {
|
|
let neural = UnifiedAttention::new("multi_head").unwrap();
|
|
assert_eq!(neural.category(), "neural");
|
|
|
|
let dag = UnifiedAttention::new("topological").unwrap();
|
|
assert_eq!(dag.category(), "dag");
|
|
|
|
let graph = UnifiedAttention::new("gat").unwrap();
|
|
assert_eq!(graph.category(), "graph");
|
|
|
|
let ssm = UnifiedAttention::new("mamba").unwrap();
|
|
assert_eq!(ssm.category(), "ssm");
|
|
}
|
|
|
|
#[wasm_bindgen_test]
|
|
fn test_softmax() {
|
|
let input = vec![1.0, 2.0, 3.0];
|
|
let output = softmax(input);
|
|
|
|
// Sum should be 1.0
|
|
let sum: f32 = output.iter().sum();
|
|
assert!((sum - 1.0).abs() < 1e-6);
|
|
|
|
// Should be monotonically increasing
|
|
assert!(output[0] < output[1]);
|
|
assert!(output[1] < output[2]);
|
|
}
|
|
|
|
#[wasm_bindgen_test]
|
|
fn test_cosine_similarity() {
|
|
let a = vec![1.0, 0.0, 0.0];
|
|
let b = vec![1.0, 0.0, 0.0];
|
|
let sim = cosine_similarity(a, b).unwrap();
|
|
assert!((sim - 1.0).abs() < 1e-6);
|
|
|
|
let c = vec![1.0, 0.0, 0.0];
|
|
let d = vec![0.0, 1.0, 0.0];
|
|
let sim2 = cosine_similarity(c, d).unwrap();
|
|
assert!(sim2.abs() < 1e-6);
|
|
}
|
|
}
|