78 lines
2.3 KiB
Rust
78 lines
2.3 KiB
Rust
//! # RTX Neural Network Library
|
|
//!
|
|
//! High-performance neural network layers and operations optimized for NVIDIA RTX GPUs.
|
|
//!
|
|
//! ## Derive Macros
|
|
//!
|
|
//! This crate re-exports derive macros for creating neural network modules:
|
|
//!
|
|
//! - `#[derive(Module)]` - Automatically implements the `Module` trait
|
|
//! - `#[derive(Config)]` - Generates builder pattern configuration structs
|
|
//!
|
|
//! See [`rtx_macros`] for detailed documentation.
|
|
|
|
#![allow(clippy::module_name_repetitions)]
|
|
|
|
pub mod functional;
|
|
pub mod init;
|
|
pub mod layers;
|
|
pub mod utils;
|
|
|
|
// Re-export derive macros for convenient use
|
|
pub use rtx_macros::{Config, Module};
|
|
|
|
// Generic module - backend-agnostic neural network layers
|
|
// Uses compile-time dispatch instead of runtime Device enum matching
|
|
pub mod generic;
|
|
pub use generic::{
|
|
GenericAdaptiveAvgPool2d, GenericAvgPool2d, GenericBatchNorm1d, GenericBatchNorm2d,
|
|
GenericConv1d, GenericConv2d, GenericDropout, GenericDropout2d, GenericDropout3d, GenericELU,
|
|
GenericEmbedding, GenericGELU, GenericGroupNorm, GenericLayerNorm, GenericLeakyReLU,
|
|
GenericLinear, GenericMLP, GenericMaxPool2d, GenericModule, GenericModule1D, GenericModule3D,
|
|
GenericModule4D, GenericMultiHeadAttention, GenericRMSNorm, GenericReLU, GenericSequential,
|
|
GenericSiLU, GenericSigmoid, GenericTanh, GenericTransformerBlock, GenericTransformerEncoder,
|
|
};
|
|
|
|
// Re-export main layer types
|
|
pub use layers::{
|
|
Module, Sequential,
|
|
activation::{GELU, ReLU},
|
|
linear::Linear,
|
|
};
|
|
|
|
/// Result type for neural network operations
|
|
pub type Result<T> = std::result::Result<T, NNError>;
|
|
|
|
/// Neural network specific errors
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum NNError {
|
|
/// Tensor operation error
|
|
#[error("Tensor error: {0}")]
|
|
Tensor(rtx_tensor::TensorError),
|
|
|
|
/// Invalid parameter error
|
|
#[error("Invalid parameter: {0}")]
|
|
InvalidParameter(String),
|
|
|
|
/// Shape mismatch error
|
|
#[error("Shape mismatch: {0}")]
|
|
ShapeMismatch(String),
|
|
|
|
/// Invalid shape error
|
|
#[error("Invalid shape: {0}")]
|
|
InvalidShape(String),
|
|
|
|
/// CUDA error
|
|
#[error("CUDA error: {0}")]
|
|
Cuda(String),
|
|
}
|
|
|
|
/// Alias for backward compatibility
|
|
pub use NNError as Error;
|
|
|
|
impl From<rtx_tensor::TensorError> for NNError {
|
|
fn from(err: rtx_tensor::TensorError) -> Self {
|
|
Self::Tensor(err)
|
|
}
|
|
}
|