tenflowers_neural/layers/
mod.rs

1//! Neural network layer implementations.
2//!
3//! This module provides a comprehensive collection of neural network layers that can be
4//! composed to build complex models. All layers implement the [`Layer`] trait, providing
5//! a consistent interface for forward and backward propagation.
6//!
7//! # Layer Categories
8//!
9//! ## Core Layers
10//!
11//! - [`Dense`]: Fully-connected (linear) layer with optional bias
12//! - [`Conv1D`], [`Conv2D`], [`Conv3D`]: Convolutional layers for 1D, 2D, and 3D data
13//! - [`ConvTranspose2D`]: Transposed convolution for upsampling
14//!
15//! ## Activation Layers
16//!
17//! - [`Activation`]: Standard activation functions (ReLU, Tanh, Sigmoid, etc.)
18//! - [`PReLU`]: Parametric ReLU with learnable parameters
19//! - [`SwiGLU`], [`GeGLU`]: Gated linear units for transformers
20//!
21//! ## Attention Mechanisms
22//!
23//! - [`MultiHeadAttention`]: Standard multi-head attention with Flash Attention support
24//! - [`MultiQueryAttention`]: Efficient multi-query attention
25//! - [`TransformerEncoder`], [`TransformerDecoder`]: Complete transformer blocks
26//! - [`FeedForwardNetwork`]: Position-wise feed-forward network
27//!
28//! ## Normalization
29//!
30//! - [`BatchNorm`]: Batch normalization with running statistics
31//! - [`LayerNorm`]: Layer normalization for transformers
32//! - [`RMSNorm`]: Root mean square normalization (LLaMA style)
33//! - [`GroupNorm`]: Group normalization for small batches
34//! - [`InstanceNorm`]: Instance normalization for style transfer
35//!
36//! ## Recurrent Layers
37//!
38//! - [`RNN`]: Basic recurrent neural network
39//! - [`LSTM`]: Long short-term memory with forget gates
40//! - [`GRU`]: Gated recurrent unit
41//!
42//! ## Regularization
43//!
44//! - [`Dropout`]: Standard dropout for regularization
45//! - [`SpatialDropout2D`]: Spatial dropout for convolutional layers
46//! - [`StochasticDepth`]: Stochastic depth (drop path)
47//!
48//! ## Pooling Operations
49//!
50//! - [`MaxPool2D`], [`AvgPool2D`]: Standard 2D pooling
51//! - [`GlobalMaxPool2D`], [`GlobalAvgPool2D`]: Global pooling
52//! - [`AdaptiveAvgPool2D`]: Adaptive pooling to target size
53//!
54//! ## Embeddings
55//!
56//! - [`Embedding`]: Token embedding layer
57//! - [`SinusoidalPositionalEncoding`]: Fixed positional encodings (Transformer)
58//! - [`LearnedPositionalEncoding`]: Learnable positional encodings
59//! - [`RotaryPositionalEmbedding`]: Rotary position embeddings (RoPE)
60//!
61//! ## Advanced Architectures
62//!
63//! - [`MambaBlock`]: Mamba state-space model block
64//! - [`StateSpaceModel`]: Generic state-space model (S4, S5)
65//! - [`MixtureOfExperts`]: Sparse mixture of experts layer
66//! - [`GraphConv`]: Graph convolution for graph neural networks
67//!
68//! # Usage Examples
69//!
70//! ## Building a Simple Network
71//!
72//! ```rust,ignore
73//! use tenflowers_neural::layers::{Dense, Activation};
74//! use tenflowers_neural::ActivationFunction;
75//! use tenflowers_core::Tensor;
76//!
77//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
78//! let layer1 = Dense::new(784, 128)?;
79//! let activation = Activation::new(ActivationFunction::ReLU);
80//! let layer2 = Dense::new(128, 10)?;
81//!
82//! let input = Tensor::zeros(&[32, 784]);
83//! let hidden = layer1.forward(&input)?;
84//! let activated = activation.forward(&hidden)?;
85//! let output = layer2.forward(&activated)?;
86//! # Ok(())
87//! # }
88//! ```
89//!
90//! ## Convolutional Network
91//!
92//! ```rust,ignore
93//! use tenflowers_neural::layers::{Conv2D, BatchNorm, MaxPool2D};
94//! use tenflowers_core::Tensor;
95//!
96//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
97//! let conv = Conv2D::new(3, 64, 3, 1, 1)?;  // in_channels, out_channels, kernel, stride, padding
98//! let bn = BatchNorm::new(64)?;
99//! let pool = MaxPool2D::new(2, 2, 0)?;      // kernel_size, stride, padding
100//!
101//! let input = Tensor::zeros(&[32, 3, 224, 224]); // NCHW format
102//! let features = conv.forward(&input)?;
103//! let normalized = bn.forward(&features, true)?; // training mode
104//! let pooled = pool.forward(&normalized)?;
105//! # Ok(())
106//! # }
107//! ```
108//!
109//! ## Transformer Block
110//!
111//! ```rust,ignore
112//! use tenflowers_neural::layers::{MultiHeadAttention, LayerNorm, FeedForwardNetwork};
113//! use tenflowers_core::Tensor;
114//!
115//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
116//! let attention = MultiHeadAttention::new(512, 8, 0.1)?; // d_model, n_heads, dropout
117//! let norm1 = LayerNorm::new(512, 1e-5)?;
118//! let ffn = FeedForwardNetwork::new(512, 2048, 0.1)?;
119//! let norm2 = LayerNorm::new(512, 1e-5)?;
120//!
121//! let x = Tensor::zeros(&[32, 128, 512]); // batch, seq_len, d_model
122//! let attn_out = attention.forward(&x, &x, &x, None)?;
123//! let x = norm1.forward(&(x.clone() + attn_out))?;
124//! let ffn_out = ffn.forward(&x)?;
125//! let output = norm2.forward(&(x + ffn_out))?;
126//! # Ok(())
127//! # }
128//! ```
129//!
130//! ## State-Space Model (Mamba)
131//!
132//! ```rust,ignore
133//! use tenflowers_neural::layers::MambaBlock;
134//! use tenflowers_core::Tensor;
135//!
136//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
137//! let mamba = MambaBlock::new(512, 16)?; // d_model, d_state
138//!
139//! let input = Tensor::zeros(&[32, 1024, 512]); // batch, seq_len, d_model
140//! let output = mamba.forward(&input)?;
141//! # Ok(())
142//! # }
143//! ```
144
145pub mod activation;
146pub mod attention;
147pub mod augmentation;
148pub mod conv;
149pub mod dense;
150pub mod dropout;
151pub mod embedding;
152pub mod gnn;
153pub mod moe;
154pub mod normalization;
155pub mod pooling;
156pub mod rnn;
157pub mod state_space;
158pub mod stochastic_depth;
159pub mod ultra_conv_simple;
160pub mod ultra_dense_simple;
161pub mod ultra_layer_manager_minimal;
162
163pub use activation::{
164    Activation, AdaptivePiecewiseLinear, AdaptivePolynomial, AdaptiveSwish, PReLU,
165    ParametricSoftplus, SwiGLU as SwiGLUActivation,
166};
167pub use attention::{
168    analyze_attention_patterns, apply_attention_mask, apply_rotary_position_embedding,
169    compute_slopes, create_causal_mask, create_padding_mask, naive_attention,
170    scaled_dot_product_attention, sinusoidal_positional_encoding, AlibiAttention, AlibiMask,
171    AlibiSlopes, AttentionStats, FeedForwardNetwork, FlashAttention, FlashConfig, GeGLU, KVCache,
172    MultiHeadAttention, MultiQueryAttention, OnlineSoftmax, RopeConfig, RopeEmbedding,
173    RotaryInterpolation, SwiGLU, TransformerDecoder, TransformerEncoder,
174};
175pub use augmentation::{CutMix, LabelSmoothing, Mixup};
176pub use conv::{Conv1D, Conv2D, Conv3D, ConvTranspose2D, DepthwiseConv2D, SeparableConv2D};
177pub use dense::Dense;
178pub use dropout::{Dropout, SpatialDropout2D};
179pub use embedding::{
180    Embedding, EmbeddingRegularization, LearnedPositionalEncoding, RotaryPositionalEmbedding,
181    SinusoidalPositionalEncoding, SparseEmbedding, SparseEmbeddingGrad,
182};
183pub use gnn::{AggregatorType, GraphAttention, GraphConv, GraphSAGE};
184pub use moe::{Expert, MixtureOfExperts, RoutingStats, TopKRouter};
185pub use normalization::{
186    BatchNorm, GroupNorm, InstanceNorm, LayerNorm, RMSNorm, SpectralNorm, WeightNorm,
187};
188pub use pooling::{
189    AdaptiveAvgPool2D, AdaptiveMaxPool2D, AvgPool2D, AvgPool3D, FractionalAvgPool2D,
190    FractionalMaxPool2D, GlobalAvgPool2D, GlobalAvgPool3D, GlobalMaxPool2D, GlobalMaxPool3D,
191    MaxPool2D, MaxPool3D, ROIAlign2D, ROIPool2D,
192};
193pub use rnn::{
194    BahdanauAttention, HierarchicalAttention, LuongAttention, LuongAttentionType,
195    ResetGateVariation, RnnNonlinearity, GRU, LSTM, RNN,
196};
197pub use state_space::{MambaBlock, StateSpaceModel};
198pub use stochastic_depth::{StochasticDepth, StochasticDepthNoResidual};
199pub use ultra_conv_simple::{ultra_conv2d, ConvPerformanceMetrics, UltraConv2D, UltraConvConfig};
200pub use ultra_dense_simple::{
201    ultra_dense, ultra_dense_no_bias, DensePerformanceMetrics, UltraDense, UltraDenseConfig,
202    UltraDenseExt,
203};
204pub use ultra_layer_manager_minimal::{
205    create_ultra_layer_manager, global_ultra_layer_manager, LayerExecutionResult, LayerId,
206    LayerMetrics, OptimizationReport, UltraLayerManager, UltraLayerManagerConfig,
207    UltraPerformanceReport,
208};
209
210use tenflowers_core::{Result, Tensor};
211
212/// Represents different types of neural network layers for ONNX export
213#[derive(Debug, Clone, PartialEq, Eq)]
214pub enum LayerType {
215    Dense,
216    Conv1D,
217    Conv2D,
218    Conv3D,
219    ConvTranspose2D,
220    MaxPool2D,
221    AvgPool2D,
222    GlobalMaxPool2D,
223    GlobalAvgPool2D,
224    BatchNorm,
225    LayerNorm,
226    RMSNorm,
227    GroupNorm,
228    Dropout,
229    LSTM,
230    GRU,
231    RNN,
232    MultiHeadAttention,
233    TransformerEncoder,
234    TransformerDecoder,
235    Embedding,
236    Activation,
237    MixtureOfExperts,
238    StateSpaceModel,
239    MambaBlock,
240    GraphConv,
241    GraphSAGE,
242    GraphAttention,
243    Unknown,
244}
245
246impl LayerType {
247    /// Convert layer type to ONNX operation type
248    pub fn to_onnx_op_type(&self) -> &'static str {
249        match self {
250            LayerType::Dense => "MatMul",
251            LayerType::Conv1D => "Conv",
252            LayerType::Conv2D => "Conv",
253            LayerType::Conv3D => "Conv",
254            LayerType::ConvTranspose2D => "ConvTranspose",
255            LayerType::MaxPool2D => "MaxPool",
256            LayerType::AvgPool2D => "AveragePool",
257            LayerType::GlobalMaxPool2D => "GlobalMaxPool",
258            LayerType::GlobalAvgPool2D => "GlobalAveragePool",
259            LayerType::BatchNorm => "BatchNormalization",
260            LayerType::LayerNorm => "LayerNormalization",
261            LayerType::RMSNorm => "LayerNormalization", // ONNX doesn't have native RMSNorm, use LayerNorm
262            LayerType::GroupNorm => "GroupNormalization",
263            LayerType::Dropout => "Dropout",
264            LayerType::LSTM => "LSTM",
265            LayerType::GRU => "GRU",
266            LayerType::RNN => "RNN",
267            LayerType::MultiHeadAttention => "MultiHeadAttention",
268            LayerType::TransformerEncoder => "Transformer",
269            LayerType::TransformerDecoder => "Transformer",
270            LayerType::Embedding => "Gather",
271            LayerType::Activation => "Relu", // Default, would need specific handling
272            LayerType::MixtureOfExperts => "Identity", // Custom operation, needs special handling
273            LayerType::StateSpaceModel => "Identity", // Custom State Space operation, needs special handling
274            LayerType::MambaBlock => "Identity", // Custom Mamba operation, needs special handling
275            LayerType::GraphConv => "Identity", // Custom Graph Convolution operation, needs special handling
276            LayerType::GraphSAGE => "Identity", // Custom GraphSAGE operation, needs special handling
277            LayerType::GraphAttention => "Identity", // Custom Graph Attention operation, needs special handling
278            LayerType::Unknown => "Identity",
279        }
280    }
281}
282
283pub trait Layer<T> {
284    fn forward(&self, input: &Tensor<T>) -> Result<Tensor<T>>;
285    fn parameters(&self) -> Vec<&Tensor<T>>;
286    fn parameters_mut(&mut self) -> Vec<&mut Tensor<T>>;
287    fn set_training(&mut self, training: bool);
288    fn clone_box(&self) -> Box<dyn Layer<T>>;
289
290    /// Forward pass for layers that accept multiple input tensors (e.g. merge/combine
291    /// layers). The default implementation only supports the single-input case: it
292    /// requires exactly one input and delegates to [`Layer::forward`], so existing
293    /// single-input layer implementations are completely unaffected by this method's
294    /// addition (no breaking change). Layers that genuinely need multiple inputs
295    /// should override this method.
296    fn forward_multi(&self, inputs: &[&Tensor<T>]) -> Result<Tensor<T>> {
297        if inputs.len() != 1 {
298            return Err(tenflowers_core::TensorError::unsupported_operation_simple(
299                format!(
300                    "this layer only supports single-input forward, but {} inputs were given \
301                 (override Layer::forward_multi to support multiple inputs)",
302                    inputs.len()
303                ),
304            ));
305        }
306        self.forward(inputs[0])
307    }
308
309    /// Returns the type of this layer for ONNX export and introspection
310    fn layer_type(&self) -> LayerType {
311        LayerType::Unknown // Default implementation
312    }
313
314    /// Set weight tensor for layers that support weights
315    /// Default implementation returns an error
316    fn set_weight(&mut self, _weight: Tensor<T>) -> Result<()> {
317        Err(tenflowers_core::TensorError::unsupported_operation_simple(
318            "This layer type does not support weight setting".to_string(),
319        ))
320    }
321
322    /// Set bias tensor for layers that support bias
323    /// Default implementation returns an error
324    fn set_bias(&mut self, _bias: Option<Tensor<T>>) -> Result<()> {
325        Err(tenflowers_core::TensorError::unsupported_operation_simple(
326            "This layer type does not support bias setting".to_string(),
327        ))
328    }
329}
tenflowers_neural/layers/mod.rs

tenflowers_neural/layers/
mod.rs
