Crate ruvector_attention

Crate ruvector_attention 

Source
Expand description

§ruvector-attention

Attention mechanisms for ruvector, including geometric, graph, and sparse attention.

This crate provides efficient implementations of various attention mechanisms:

  • Scaled dot-product attention
  • Multi-head attention with parallel processing
  • Graph attention for GNN applications
  • Geometric attention in hyperbolic spaces
  • Sparse attention patterns

§Features

  • SIMD Acceleration: Optional SIMD optimizations for performance
  • Parallel Processing: Rayon-based parallel head computation
  • WASM Support: WebAssembly compilation support
  • NAPI Bindings: Node.js bindings for JavaScript integration

§Example

use ruvector_attention::{
    attention::ScaledDotProductAttention,
    traits::Attention,
};

// Create scaled dot-product attention
let attention = ScaledDotProductAttention::new(512);

// Prepare inputs
let query = vec![1.0; 512];
let keys = vec![vec![0.5; 512], vec![0.3; 512]];
let values = vec![vec![1.0; 512], vec![2.0; 512]];

let keys_refs: Vec<&[f32]> = keys.iter().map(|k| k.as_slice()).collect();
let values_refs: Vec<&[f32]> = values.iter().map(|v| v.as_slice()).collect();

// Compute attention
let output = attention.compute(&query, &keys_refs, &values_refs).unwrap();
assert_eq!(output.len(), 512);

Re-exports§

pub use attention::MultiHeadAttention;
pub use attention::ScaledDotProductAttention;
pub use config::AttentionConfig;
pub use config::GraphAttentionConfig;
pub use config::SparseAttentionConfig;
pub use error::AttentionError;
pub use error::AttentionResult;
pub use hyperbolic::exp_map;
pub use hyperbolic::log_map;
pub use hyperbolic::mobius_add;
pub use hyperbolic::poincare_distance;
pub use hyperbolic::project_to_ball;
pub use hyperbolic::HyperbolicAttention;
pub use hyperbolic::HyperbolicAttentionConfig;
pub use hyperbolic::MixedCurvatureAttention;
pub use hyperbolic::MixedCurvatureConfig;
pub use traits::Attention;
pub use traits::EdgeInfo;
pub use traits::GeometricAttention;
pub use traits::Gradients;
pub use traits::GraphAttention;
pub use traits::SparseAttention;
pub use traits::SparseMask;
pub use traits::TrainableAttention;
pub use sparse::AttentionMask;
pub use sparse::FlashAttention;
pub use sparse::LinearAttention;
pub use sparse::LocalGlobalAttention;
pub use sparse::SparseMaskBuilder;
pub use moe::Expert;
pub use moe::ExpertType;
pub use moe::HyperbolicExpert;
pub use moe::LearnedRouter;
pub use moe::LinearExpert;
pub use moe::MoEAttention;
pub use moe::MoEConfig;
pub use moe::Router;
pub use moe::StandardExpert;
pub use moe::TopKRouting;
pub use graph::DualSpaceAttention;
pub use graph::DualSpaceConfig;
pub use graph::EdgeFeaturedAttention;
pub use graph::EdgeFeaturedConfig;
pub use graph::GraphRoPE;
pub use graph::RoPEConfig;
pub use training::Adam;
pub use training::AdamW;
pub use training::CurriculumScheduler;
pub use training::CurriculumStage;
pub use training::DecayType;
pub use training::HardNegativeMiner;
pub use training::InfoNCELoss;
pub use training::LocalContrastiveLoss;
pub use training::Loss;
pub use training::MiningStrategy;
pub use training::NegativeMiner;
pub use training::Optimizer;
pub use training::Reduction;
pub use training::SpectralRegularization;
pub use training::TemperatureAnnealing;
pub use training::SGD;
pub use sdk::presets;
pub use sdk::AttentionBuilder;
pub use sdk::AttentionPipeline;
pub use transport::CentroidCache;
pub use transport::CentroidOTAttention;
pub use transport::CentroidOTConfig;
pub use transport::ProjectionCache;
pub use transport::SlicedWassersteinAttention;
pub use transport::SlicedWassersteinConfig;
pub use transport::WindowCache;
pub use curvature::ComponentQuantizer;
pub use curvature::FusedCurvatureConfig;
pub use curvature::MixedCurvatureCache;
pub use curvature::MixedCurvatureFusedAttention;
pub use curvature::QuantizationConfig;
pub use curvature::QuantizedVector;
pub use curvature::TangentSpaceMapper;
pub use curvature::TangentSpaceConfig;
pub use topology::AttentionMode;
pub use topology::AttentionPolicy;
pub use topology::CoherenceMetric;
pub use topology::PolicyConfig;
pub use topology::TopologyGatedAttention;
pub use topology::TopologyGatedConfig;
pub use topology::WindowCoherence;
pub use info_geometry::FisherConfig;
pub use info_geometry::FisherMetric;
pub use info_geometry::NaturalGradient;
pub use info_geometry::NaturalGradientConfig;
pub use info_bottleneck::DiagonalGaussian;
pub use info_bottleneck::IBConfig;
pub use info_bottleneck::InformationBottleneck;
pub use info_bottleneck::KLDivergence;
pub use pde_attention::DiffusionAttention;
pub use pde_attention::DiffusionConfig;
pub use pde_attention::GraphLaplacian;
pub use pde_attention::LaplacianType;
pub use unified_report::AttentionRecommendation;
pub use unified_report::GeometryReport;
pub use unified_report::MetricType;
pub use unified_report::MetricValue;
pub use unified_report::ReportBuilder;
pub use unified_report::ReportConfig;

Modules§

attention
Attention mechanism implementations.
config
Configuration types for attention mechanisms.
curvature
Mixed Curvature Attention
error
Error types for the ruvector-attention crate.
graph
Graph attention mechanisms for GNN applications
hyperbolic
Hyperbolic Attention Module
info_bottleneck
Information Bottleneck
info_geometry
Information Geometry for Attention
moe
Mixture of Experts (MoE) attention mechanisms
pde_attention
PDE-Based Attention
sdk
ruvector-attention SDK
sparse
Sparse attention mechanisms for efficient computation on long sequences
topology
Topology Gated Attention
training
Training utilities for attention-based graph neural networks
traits
Trait definitions for attention mechanisms.
transport
Optimal Transport Attention
unified_report
Unified Geometry Report
utils
Utility functions for attention mechanisms.

Constants§

VERSION
Library version