scivex-nn 0.1.0

#![allow(
    clippy::missing_errors_doc,
    clippy::missing_panics_doc,
    clippy::must_use_candidate,
    clippy::return_self_not_must_use,
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss,
    clippy::cast_lossless,
    clippy::many_single_char_names,
    clippy::similar_names,
    clippy::doc_markdown,
    clippy::module_name_repetitions
)]
//! # scivex-nn
//!
//! Neural networks with reverse-mode automatic differentiation for the
//! Scivex ecosystem.
//!
//! ## Modules
//!
//! | Module | Contents |
//! |--------|----------|
//! | [`variable`] | `Variable<T>` — autograd computation graph node |
//! | [`ops`] | Differentiable operations (add, mul, matmul, etc.) |
//! | [`functional`] | Activation functions (relu, sigmoid, tanh, softmax) |
//! | [`layer`] | `Layer` trait, `Linear`, `Conv1d/2d/3d`, `BatchNorm1d/2d`, pooling, RNN/LSTM/GRU, attention, `Sequential` |
//! | [`optim`] | `Optimizer` trait, `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, LR schedulers |
//! | [`loss`] | Loss functions (MSE, cross-entropy, BCE, Huber, focal, KL, hinge, smooth L1) |
//! | [`init`] | Weight initialization (Xavier, Kaiming) |
//! | [`persist`] | `save_weights`, `load_weights` — binary weight persistence |
//! | [`data`] | `Dataset` trait, `TensorDataset`, `DataLoader` |

/// Dataset, DataLoader, and batching utilities.
pub mod data;
/// Neural network error types.
pub mod error;
/// Model interpretability: Integrated Gradients, SmoothGrad.
pub mod explain;
/// Activation functions (ReLU, sigmoid, tanh, softmax).
pub mod functional;
/// Weight initialization strategies (Xavier, Kaiming).
pub mod init;
/// Layer trait and built-in layers (Linear, BatchNorm, Dropout).
pub mod layer;
/// Loss functions (MSE, cross-entropy, BCE).
pub mod loss;
/// ONNX model loading and inference.
pub mod onnx;
/// Differentiable tensor operations for the autograd graph.
pub mod ops;
/// Optimizers (SGD, Adam, AdamW, RMSprop, Adagrad).
pub mod optim;
/// Weight persistence: save and load model parameters.
pub mod persist;
/// Model serialization formats (SafeTensors, GGUF).
pub mod serialize;
/// In-process model inference server with batching support.
pub mod serve;
/// Training utilities (Trainer, callbacks, gradient clipping).
pub mod training;
/// Autograd computation graph node.
pub mod variable;

/// GPU-accelerated training backend.
#[cfg(feature = "gpu")]
pub mod gpu;

pub use error::{NnError, Result};
pub use variable::Variable;

/// Convenience re-exports.
pub mod prelude {
    pub use crate::data::{DataLoader, Dataset, TensorDataset};
    pub use crate::error::{NnError, Result};
    pub use crate::explain::{
        IntegratedGradientsResult, integrated_gradients, integrated_gradients_zero_baseline,
        smooth_gradients,
    };
    pub use crate::functional::{log_softmax, relu, sigmoid, softmax, tanh_fn};
    pub use crate::init::{kaiming_normal, kaiming_uniform, xavier_normal, xavier_uniform};
    pub use crate::layer::{
        AvgPool1d, AvgPool2d, BatchNorm1d, BatchNorm2d, Conv1d, Conv2d, Conv3d, Dropout, Embedding,
        FlashAttention, Flatten, GATConv, GCNConv, GRU, GroupedQueryAttention, LSTM, Layer,
        LayerNorm, Linear, MaxPool1d, MaxPool2d, MultiHeadAttention, MultiQueryAttention, ReLU,
        RotaryPositionalEncoding, SAGEConv, Sequential, Sigmoid, SimpleRNN,
        SinusoidalPositionalEncoding, Tanh, TransformerDecoderLayer, TransformerEncoderLayer,
        causal_mask,
    };
    pub use crate::loss::{
        bce_loss, cross_entropy_loss, focal_loss, hinge_loss, huber_loss, kl_divergence, mse_loss,
        smooth_l1_loss,
    };
    pub use crate::onnx::{
        OnnxAttribute, OnnxAttributeValue, OnnxDataType, OnnxGraph, OnnxInferenceSession,
        OnnxModel, OnnxNode, OnnxOpsetImport, OnnxTensor, OnnxValueInfo, load_onnx,
    };
    pub use crate::ops::{add, add_bias, matmul, mean, mul, neg, pow, scalar_mul, sub, sum};
    pub use crate::optim::{
        Adagrad, Adam, AdamW, CosineAnnealingLR, ExponentialLR, LinearLR, LrScheduler, Optimizer,
        RMSprop, ReduceLROnPlateau, SGD, StepLR, WarmupCosineDecay,
    };
    pub use crate::persist::{load_weights, save_weights};
    pub use crate::serialize::{
        GgufFile, GgufValue, load_gguf, load_safetensors, save_gguf, save_safetensors,
    };
    pub use crate::serve::{
        FnModel, InferenceConfig, InferenceModel, InferenceRequest, InferenceResponse,
        InferenceServer, InferenceStats,
    };
    pub use crate::training::{
        AmpConfig, Callback, CallbackAction, EarlyStopping, GradAccumulator, GradScaler,
        LossLogger, LrFinder, LrFinderResult, ModelCheckpoint, Trainer, TrainingHistory,
        cast_params, cast_variable, clip_grad_norm, clip_grad_value,
    };
    pub use crate::variable::Variable;

    #[cfg(feature = "gpu")]
    pub use crate::gpu::prelude::*;
}