Skip to main content

ruvector_gnn/
lib.rs

1//! # RuVector GNN
2//!
3//! Graph Neural Network capabilities for RuVector, providing tensor operations,
4//! GNN layers, compression, and differentiable search.
5//!
6//! ## Forgetting Mitigation (Issue #17)
7//!
8//! This crate includes comprehensive forgetting mitigation for continual learning:
9//!
10//! - **Adam Optimizer**: Full implementation with momentum and bias correction
11//! - **Replay Buffer**: Experience replay with reservoir sampling for uniform coverage
12//! - **EWC (Elastic Weight Consolidation)**: Prevents catastrophic forgetting
13//! - **Learning Rate Scheduling**: Multiple strategies including warmup and plateau detection
14//!
15//! ### Usage Example
16//!
17//! ```rust,ignore
18//! use ruvector_gnn::{
19//!     training::{Optimizer, OptimizerType},
20//!     replay::ReplayBuffer,
21//!     ewc::ElasticWeightConsolidation,
22//!     scheduler::{LearningRateScheduler, SchedulerType},
23//! };
24//!
25//! // Create Adam optimizer
26//! let mut optimizer = Optimizer::new(OptimizerType::Adam {
27//!     learning_rate: 0.001,
28//!     beta1: 0.9,
29//!     beta2: 0.999,
30//!     epsilon: 1e-8,
31//! });
32//!
33//! // Create replay buffer for experience replay
34//! let mut replay = ReplayBuffer::new(10000);
35//!
36//! // Create EWC for preventing forgetting
37//! let mut ewc = ElasticWeightConsolidation::new(0.4);
38//!
39//! // Create learning rate scheduler
40//! let mut scheduler = LearningRateScheduler::new(
41//!     SchedulerType::CosineAnnealing { t_max: 100, eta_min: 1e-6 },
42//!     0.001
43//! );
44//! ```
45
46// Public-API doc coverage tracked separately; allow incomplete field-level docs for now.
47#![allow(missing_docs)]
48#![deny(unsafe_op_in_unsafe_fn)]
49
50pub mod compress;
51pub mod error;
52pub mod ewc;
53pub mod graphmae;
54pub mod layer;
55pub mod query;
56pub mod replay;
57pub mod scheduler;
58pub mod search;
59pub mod tensor;
60pub mod training;
61
62#[cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
63pub mod mmap;
64
65#[cfg(all(feature = "cold-tier", not(target_arch = "wasm32")))]
66pub mod cold_tier;
67
68// Re-export commonly used types
69pub use compress::{CompressedTensor, CompressionLevel, TensorCompress};
70pub use error::{GnnError, Result};
71pub use ewc::ElasticWeightConsolidation;
72pub use graphmae::{
73    mse_loss, sce_loss, FeatureMasking, GATEncoder, GraphData, GraphMAE, GraphMAEConfig,
74    GraphMAEDecoder, LossFn, MaskResult,
75};
76pub use layer::RuvectorLayer;
77pub use query::{QueryMode, QueryResult, RuvectorQuery, SubGraph};
78pub use replay::{DistributionStats, ReplayBuffer, ReplayEntry};
79pub use scheduler::{LearningRateScheduler, SchedulerType};
80pub use search::{cosine_similarity, differentiable_search, hierarchical_forward};
81pub use training::{
82    info_nce_loss, local_contrastive_loss, sgd_step, Loss, LossType, OnlineConfig, Optimizer,
83    OptimizerType, TrainConfig,
84};
85
86#[cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
87pub use mmap::{AtomicBitmap, MmapGradientAccumulator, MmapManager};
88
89#[cfg(test)]
90mod tests {
91    #[test]
92    fn test_basic() {
93        // Basic smoke test to ensure the crate compiles
94    }
95}