scivex_nn/
lib.rs

1#![allow(
2    clippy::missing_errors_doc,
3    clippy::missing_panics_doc,
4    clippy::must_use_candidate,
5    clippy::return_self_not_must_use,
6    clippy::cast_precision_loss,
7    clippy::cast_possible_truncation,
8    clippy::cast_sign_loss,
9    clippy::cast_lossless,
10    clippy::many_single_char_names,
11    clippy::similar_names,
12    clippy::doc_markdown,
13    clippy::module_name_repetitions
14)]
15//! # scivex-nn
16//!
17//! Neural networks with reverse-mode automatic differentiation for the
18//! Scivex ecosystem.
19//!
20//! ## Modules
21//!
22//! | Module | Contents |
23//! |--------|----------|
24//! | [`variable`] | `Variable<T>` — autograd computation graph node |
25//! | [`ops`] | Differentiable operations (add, mul, matmul, etc.) |
26//! | [`functional`] | Activation functions (relu, sigmoid, tanh, softmax) |
27//! | [`layer`] | `Layer` trait, `Linear`, `Conv1d/2d/3d`, `BatchNorm1d/2d`, pooling, RNN/LSTM/GRU, attention, `Sequential` |
28//! | [`optim`] | `Optimizer` trait, `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, LR schedulers |
29//! | [`loss`] | Loss functions (MSE, cross-entropy, BCE, Huber, focal, KL, hinge, smooth L1) |
30//! | [`init`] | Weight initialization (Xavier, Kaiming) |
31//! | [`persist`] | `save_weights`, `load_weights` — binary weight persistence |
32//! | [`data`] | `Dataset` trait, `TensorDataset`, `DataLoader` |
33
34/// Dataset, DataLoader, and batching utilities.
35pub mod data;
36/// Neural network error types.
37pub mod error;
38/// Model interpretability: Integrated Gradients, SmoothGrad.
39pub mod explain;
40/// Activation functions (ReLU, sigmoid, tanh, softmax).
41pub mod functional;
42/// Weight initialization strategies (Xavier, Kaiming).
43pub mod init;
44/// Layer trait and built-in layers (Linear, BatchNorm, Dropout).
45pub mod layer;
46/// Loss functions (MSE, cross-entropy, BCE).
47pub mod loss;
48/// ONNX model loading and inference.
49pub mod onnx;
50/// Differentiable tensor operations for the autograd graph.
51pub mod ops;
52/// Optimizers (SGD, Adam, AdamW, RMSprop, Adagrad).
53pub mod optim;
54/// Weight persistence: save and load model parameters.
55pub mod persist;
56/// Model serialization formats (SafeTensors, GGUF).
57pub mod serialize;
58/// In-process model inference server with batching support.
59pub mod serve;
60/// Training utilities (Trainer, callbacks, gradient clipping).
61pub mod training;
62/// Autograd computation graph node.
63pub mod variable;
64
65/// GPU-accelerated training backend.
66#[cfg(feature = "gpu")]
67pub mod gpu;
68
69pub use error::{NnError, Result};
70pub use variable::Variable;
71
72/// Convenience re-exports.
73pub mod prelude {
74    pub use crate::data::{DataLoader, Dataset, TensorDataset};
75    pub use crate::error::{NnError, Result};
76    pub use crate::explain::{
77        IntegratedGradientsResult, integrated_gradients, integrated_gradients_zero_baseline,
78        smooth_gradients,
79    };
80    pub use crate::functional::{log_softmax, relu, sigmoid, softmax, tanh_fn};
81    pub use crate::init::{kaiming_normal, kaiming_uniform, xavier_normal, xavier_uniform};
82    pub use crate::layer::{
83        AvgPool1d, AvgPool2d, BatchNorm1d, BatchNorm2d, Conv1d, Conv2d, Conv3d, Dropout, Embedding,
84        FlashAttention, Flatten, GATConv, GCNConv, GRU, GroupedQueryAttention, LSTM, Layer,
85        LayerNorm, Linear, MaxPool1d, MaxPool2d, MultiHeadAttention, MultiQueryAttention, ReLU,
86        RotaryPositionalEncoding, SAGEConv, Sequential, Sigmoid, SimpleRNN,
87        SinusoidalPositionalEncoding, Tanh, TransformerDecoderLayer, TransformerEncoderLayer,
88        causal_mask,
89    };
90    pub use crate::loss::{
91        bce_loss, cross_entropy_loss, focal_loss, hinge_loss, huber_loss, kl_divergence, mse_loss,
92        smooth_l1_loss,
93    };
94    pub use crate::onnx::{
95        OnnxAttribute, OnnxAttributeValue, OnnxDataType, OnnxGraph, OnnxInferenceSession,
96        OnnxModel, OnnxNode, OnnxOpsetImport, OnnxTensor, OnnxValueInfo, load_onnx,
97    };
98    pub use crate::ops::{add, add_bias, matmul, mean, mul, neg, pow, scalar_mul, sub, sum};
99    pub use crate::optim::{
100        Adagrad, Adam, AdamW, CosineAnnealingLR, ExponentialLR, LinearLR, LrScheduler, Optimizer,
101        RMSprop, ReduceLROnPlateau, SGD, StepLR, WarmupCosineDecay,
102    };
103    pub use crate::persist::{load_weights, save_weights};
104    pub use crate::serialize::{
105        GgufFile, GgufValue, load_gguf, load_safetensors, save_gguf, save_safetensors,
106    };
107    pub use crate::serve::{
108        FnModel, InferenceConfig, InferenceModel, InferenceRequest, InferenceResponse,
109        InferenceServer, InferenceStats,
110    };
111    pub use crate::training::{
112        AmpConfig, Callback, CallbackAction, EarlyStopping, GradAccumulator, GradScaler,
113        LossLogger, LrFinder, LrFinderResult, ModelCheckpoint, Trainer, TrainingHistory,
114        cast_params, cast_variable, clip_grad_norm, clip_grad_value,
115    };
116    pub use crate::variable::Variable;
117
118    #[cfg(feature = "gpu")]
119    pub use crate::gpu::prelude::*;
120}
scivex_nn/lib.rs

scivex_nn/
lib.rs
