trustformers_core/traits.rs
1//! Core traits defining the fundamental abstractions of TrustformeRS.
2//!
3//! This module contains the essential traits that form the foundation of the TrustformeRS
4//! transformer library. These traits define the interfaces for models, layers, configuration,
5//! tokenization, optimization, and parameter initialization.
6//!
7//! # Overview
8//!
9//! The traits in this module establish a consistent API across all transformer implementations:
10//!
11//! - [`Model`]: The main trait for transformer models with forward pass and loading capabilities
12//! - [`Layer`]: Building blocks for neural network architectures
13//! - [`Config`]: Configuration management for models and components
14//! - [`WeightReader`]: Interface for loading pretrained model weights
15//! - [`Tokenizer`]: Text tokenization and encoding/decoding
16//! - [`Optimizer`]: Parameter optimization algorithms
17//! - [`ParameterInit`]: Weight initialization strategies
18//!
19//! # Examples
20//!
21//! ```no_run
22//! use trustformers_core::traits::{Model, Config};
23//! use trustformers_core::tensor::Tensor;
24//! use trustformers_core::errors::Result;
25//! use std::io::Read;
26//! use serde::{Deserialize, Serialize};
27//!
28//! // Example model configuration
29//! #[derive(Debug, Deserialize, Serialize)]
30//! struct MyConfig { hidden_size: usize }
31//! impl Config for MyConfig {
32//! fn architecture(&self) -> &'static str { "my_model" }
33//! }
34//!
35//! // Example model implementation
36//! struct MyModel {
37//! config: MyConfig,
38//! // ... model layers
39//! }
40//!
41//! impl Model for MyModel {
42//! type Config = MyConfig;
43//! type Input = Tensor;
44//! type Output = Tensor;
45//!
46//! fn forward(&self, input: Self::Input) -> Result<Self::Output> {
47//! // Model forward pass implementation
48//! Ok(input)
49//! }
50//!
51//! fn load_pretrained(&mut self, _reader: &mut dyn Read) -> Result<()> {
52//! // Load pretrained weights
53//! Ok(())
54//! }
55//!
56//! fn get_config(&self) -> &Self::Config {
57//! &self.config
58//! }
59//!
60//! fn num_parameters(&self) -> usize { 0 }
61//! }
62//! ```
63
64use crate::errors::Result;
65use crate::tensor::Tensor;
66use serde::{Deserialize, Serialize};
67use std::io::Read;
68
69/// The main trait for transformer models.
70///
71/// This trait defines the interface that all transformer models must implement,
72/// providing a consistent API for forward passes, weight loading, and configuration access.
73///
74/// # Type Parameters
75///
76/// - `Config`: The configuration type for this model, must implement [`Config`]
77/// - `Input`: The input type for the model's forward pass
78/// - `Output`: The output type produced by the model
79///
80/// # Thread Safety
81///
82/// Models must be `Send + Sync` to support multi-threaded inference and training.
83///
84/// # Example
85///
86/// ```no_run
87/// use trustformers_core::traits::{Model, Config};
88/// use trustformers_core::tensor::Tensor;
89/// use trustformers_core::errors::Result;
90/// use std::io::Read;
91/// use serde::{Deserialize, Serialize};
92///
93/// #[derive(Deserialize, Serialize)]
94/// struct BertConfig {
95/// hidden_size: usize,
96/// num_attention_heads: usize,
97/// // ... other config fields
98/// }
99///
100/// impl Config for BertConfig {
101/// fn architecture(&self) -> &'static str {
102/// "bert"
103/// }
104/// }
105///
106/// struct BertModel {
107/// config: BertConfig,
108/// // ... model layers
109/// }
110///
111/// impl Model for BertModel {
112/// type Config = BertConfig;
113/// type Input = Tensor;
114/// type Output = Tensor;
115///
116/// fn forward(&self, input: Self::Input) -> Result<Self::Output> {
117/// // BERT forward pass implementation
118/// Ok(input)
119/// }
120///
121/// fn load_pretrained(&mut self, _reader: &mut dyn Read) -> Result<()> {
122/// // Load BERT weights from reader
123/// Ok(())
124/// }
125///
126/// fn get_config(&self) -> &Self::Config {
127/// &self.config
128/// }
129///
130/// fn num_parameters(&self) -> usize { 0 }
131/// }
132/// ```
133pub trait Model: Send + Sync {
134 type Config: Config;
135 type Input;
136 type Output;
137
138 /// Performs a forward pass through the model.
139 ///
140 /// # Arguments
141 ///
142 /// * `input` - The input data for the model
143 ///
144 /// # Returns
145 ///
146 /// Returns `Ok(output)` on success, or an error if the forward pass fails.
147 ///
148 /// # Errors
149 ///
150 /// May return errors for:
151 /// - Invalid input dimensions
152 /// - Numerical computation errors
153 /// - Out of memory conditions
154 fn forward(&self, input: Self::Input) -> Result<Self::Output>;
155
156 /// Loads pretrained weights into the model.
157 ///
158 /// This method reads model weights from a reader (typically a file or network stream)
159 /// and updates the model's parameters accordingly.
160 ///
161 /// # Arguments
162 ///
163 /// * `reader` - A reader providing access to the pretrained weight data
164 ///
165 /// # Returns
166 ///
167 /// Returns `Ok(())` on successful loading, or an error if loading fails.
168 ///
169 /// # Errors
170 ///
171 /// May return errors for:
172 /// - IO errors while reading
173 /// - Incompatible weight formats
174 /// - Mismatched tensor shapes
175 fn load_pretrained(&mut self, reader: &mut dyn Read) -> Result<()>;
176
177 /// Returns a reference to the model's configuration.
178 ///
179 /// # Returns
180 ///
181 /// A reference to the model's configuration object.
182 fn get_config(&self) -> &Self::Config;
183
184 /// Returns the total number of parameters in the model.
185 ///
186 /// This method calculates the total number of trainable parameters
187 /// across all layers in the model. It's useful for compression metrics,
188 /// model size analysis, and memory usage calculations.
189 ///
190 /// # Returns
191 ///
192 /// The total number of parameters as a `usize`.
193 ///
194 /// # Example
195 ///
196 /// ```no_run
197 /// use trustformers_core::traits::Model;
198 ///
199 /// fn analyze_model_size<M: Model>(model: &M) {
200 /// let params = model.num_parameters();
201 /// let memory_mb = params * 4 / (1024 * 1024); // Assuming f32 weights
202 /// println!("Model has {} parameters ({} MB)", params, memory_mb);
203 /// }
204 /// ```
205 fn num_parameters(&self) -> usize;
206}
207
208/// A building block for neural network architectures.
209///
210/// The `Layer` trait represents a single computational unit in a neural network,
211/// such as a linear transformation, attention mechanism, or normalization layer.
212/// Layers can be composed together to build complete models.
213///
214/// # Type Parameters
215///
216/// - `Input`: The input type accepted by this layer
217/// - `Output`: The output type produced by this layer
218///
219/// # Thread Safety
220///
221/// Layers must be `Send + Sync` to support parallel computation.
222///
223/// # Example
224///
225/// ```no_run
226/// use trustformers_core::traits::Layer;
227/// use trustformers_core::tensor::Tensor;
228/// use trustformers_core::errors::Result;
229///
230/// struct LinearLayer {
231/// weight: Tensor,
232/// bias: Option<Tensor>,
233/// }
234///
235/// impl Layer for LinearLayer {
236/// type Input = Tensor;
237/// type Output = Tensor;
238///
239/// fn forward(&self, input: Self::Input) -> Result<Self::Output> {
240/// // Compute linear transformation: y = xW^T + b
241/// let output = input.matmul(&self.weight.transpose(0, 1)?)?;
242/// if let Some(bias) = &self.bias {
243/// output.add(bias)
244/// } else {
245/// Ok(output)
246/// }
247/// }
248/// }
249/// ```
250pub trait Layer: Send + Sync {
251 type Input;
252 type Output;
253
254 /// Performs the forward computation of this layer.
255 ///
256 /// # Arguments
257 ///
258 /// * `input` - The input data to process
259 ///
260 /// # Returns
261 ///
262 /// Returns `Ok(output)` containing the layer's output, or an error if computation fails.
263 ///
264 /// # Errors
265 ///
266 /// May return errors for:
267 /// - Invalid input dimensions
268 /// - Numerical errors during computation
269 /// - Resource allocation failures
270 fn forward(&self, input: Self::Input) -> Result<Self::Output>;
271}
272
273/// Configuration trait for models and components.
274///
275/// This trait provides a standardized interface for configuration objects
276/// that can be serialized, deserialized, and validated. All model configurations
277/// must implement this trait to ensure compatibility with the TrustformeRS ecosystem.
278///
279/// # Requirements
280///
281/// Implementing types must be serializable and deserializable using serde.
282///
283/// # Example
284///
285/// ```no_run
286/// use trustformers_core::traits::Config;
287/// use trustformers_core::errors::{Result, TrustformersError};
288/// use serde::{Deserialize, Serialize};
289///
290/// #[derive(Debug, Clone, Deserialize, Serialize)]
291/// struct GPT2Config {
292/// vocab_size: usize,
293/// hidden_size: usize,
294/// num_layers: usize,
295/// num_heads: usize,
296/// }
297///
298/// impl Config for GPT2Config {
299/// fn validate(&self) -> Result<()> {
300/// if self.hidden_size % self.num_heads != 0 {
301/// return Err(TrustformersError::invalid_input(
302/// "hidden_size must be divisible by num_heads".to_string(),
303/// ));
304/// }
305/// Ok(())
306/// }
307///
308/// fn architecture(&self) -> &'static str {
309/// "gpt2"
310/// }
311/// }
312/// ```
313pub trait Config: for<'de> Deserialize<'de> + Serialize {
314 /// Validates the configuration for correctness.
315 ///
316 /// This method should check that all configuration parameters are valid
317 /// and compatible with each other. The default implementation accepts
318 /// all configurations as valid.
319 ///
320 /// # Returns
321 ///
322 /// Returns `Ok(())` if the configuration is valid, or an error describing
323 /// the validation failure.
324 ///
325 /// # Example
326 ///
327 /// Common validations include:
328 /// - Checking that dimensions are compatible
329 /// - Verifying that values are within acceptable ranges
330 /// - Ensuring required fields are properly set
331 fn validate(&self) -> Result<()> {
332 Ok(())
333 }
334
335 /// Returns the architecture name for this configuration.
336 ///
337 /// This should return a static string identifying the model architecture,
338 /// such as "bert", "gpt2", "t5", etc. This is used for model registration
339 /// and automatic model selection.
340 ///
341 /// # Returns
342 ///
343 /// A static string slice containing the architecture name.
344 fn architecture(&self) -> &'static str;
345}
346
347/// Interface for reading model weights from various sources.
348///
349/// `WeightReader` provides an abstraction over different weight storage formats,
350/// allowing models to load pretrained parameters from files, network sources,
351/// or other storage backends.
352///
353/// # Supported Formats
354///
355/// Implementations may support various formats including:
356/// - SafeTensors (.safetensors)
357/// - PyTorch checkpoints (.pt, .bin)
358/// - NumPy arrays (.npz)
359/// - Custom formats
360///
361/// # Example
362///
363/// ```no_run
364/// use trustformers_core::traits::WeightReader;
365/// use trustformers_core::tensor::Tensor;
366/// use trustformers_core::errors::Result;
367///
368/// struct SafeTensorsReader {
369/// // ... implementation details
370/// }
371///
372/// impl WeightReader for SafeTensorsReader {
373/// fn read_tensor(&mut self, _name: &str) -> Result<Tensor> {
374/// // Read tensor from SafeTensors file
375/// Ok(Tensor::zeros(&[768, 768])?)
376/// }
377///
378/// fn list_tensors(&self) -> Vec<String> {
379/// vec![
380/// "bert.embeddings.word_embeddings.weight".to_string(),
381/// "bert.encoder.layer.0.attention.self.query.weight".to_string(),
382/// // ... more tensor names
383/// ]
384/// }
385/// }
386/// ```
387pub trait WeightReader {
388 /// Reads a tensor by name from the weight source.
389 ///
390 /// # Arguments
391 ///
392 /// * `name` - The name/key of the tensor to read (e.g., "encoder.layer.0.weight")
393 ///
394 /// # Returns
395 ///
396 /// Returns `Ok(tensor)` containing the requested tensor, or an error if the
397 /// tensor cannot be found or loaded.
398 ///
399 /// # Errors
400 ///
401 /// May return errors for:
402 /// - Tensor not found with the given name
403 /// - IO errors while reading
404 /// - Corrupted or invalid tensor data
405 /// - Unsupported tensor format
406 fn read_tensor(&mut self, name: &str) -> Result<Tensor>;
407
408 /// Lists all available tensor names in the weight source.
409 ///
410 /// This method is useful for debugging and for discovering the structure
411 /// of saved model weights.
412 ///
413 /// # Returns
414 ///
415 /// A vector containing the names of all available tensors.
416 fn list_tensors(&self) -> Vec<String>;
417}
418
419/// Text tokenization interface for transformer models.
420///
421/// The `Tokenizer` trait provides methods for converting between text and token IDs,
422/// which is essential for preparing input data for transformer models. Implementations
423/// may use various tokenization algorithms such as WordPiece, BPE, or SentencePiece.
424///
425/// # Thread Safety
426///
427/// Tokenizers must be `Send + Sync` to support concurrent tokenization.
428///
429/// # Example
430///
431/// ```no_run
432/// use trustformers_core::traits::{Tokenizer, TokenizedInput};
433/// use trustformers_core::errors::Result;
434/// use std::collections::HashMap;
435///
436/// struct BertTokenizer {
437/// vocab: HashMap<String, u32>,
438/// // ... other fields
439/// }
440///
441/// impl Tokenizer for BertTokenizer {
442/// fn encode(&self, _text: &str) -> Result<TokenizedInput> {
443/// // Tokenize text into subwords
444/// let tokens: Vec<u32> = vec![101, 2023, 2003, 1037, 3231, 102]; // [CLS] this is a test [SEP]
445/// Ok(TokenizedInput::new(tokens, vec![1; 6]))
446/// }
447///
448/// fn encode_pair(&self, text: &str, text2: &str) -> Result<TokenizedInput> {
449/// // Encode two texts for tasks like question answering
450/// let tokens1 = self.encode(text)?;
451/// let tokens2 = self.encode(text2)?;
452/// // Combine tokens with separator
453/// let len1 = tokens1.input_ids.len();
454/// let len2 = tokens2.input_ids.len();
455/// let mut combined_ids = tokens1.input_ids;
456/// combined_ids.extend_from_slice(&tokens2.input_ids);
457/// let combined_len = combined_ids.len();
458/// Ok(TokenizedInput::with_token_type_ids(
459/// combined_ids,
460/// vec![1; combined_len],
461/// Some(vec![0; len1].into_iter().chain(vec![1; len2]).collect()),
462/// ))
463/// }
464///
465/// fn decode(&self, _ids: &[u32]) -> Result<String> {
466/// // Convert token IDs back to text
467/// Ok("this is a test".to_string())
468/// }
469///
470/// fn vocab_size(&self) -> usize {
471/// 30522 // BERT base vocabulary size
472/// }
473///
474/// fn get_vocab(&self) -> HashMap<String, u32> {
475/// self.vocab.clone()
476/// }
477///
478/// fn token_to_id(&self, token: &str) -> Option<u32> {
479/// self.vocab.get(token).copied()
480/// }
481///
482/// fn id_to_token(&self, id: u32) -> Option<String> {
483/// self.vocab.iter().find(|(_, &v)| v == id).map(|(k, _)| k.clone())
484/// }
485/// }
486/// ```
487pub trait Tokenizer: Send + Sync {
488 /// Encodes a single text string into tokens.
489 ///
490 /// # Arguments
491 ///
492 /// * `text` - The input text to tokenize
493 ///
494 /// # Returns
495 ///
496 /// Returns a `TokenizedInput` containing:
497 /// - `input_ids`: The token IDs
498 /// - `attention_mask`: Binary mask indicating real vs padding tokens
499 /// - `token_type_ids`: Optional segment IDs for models like BERT
500 ///
501 /// # Errors
502 ///
503 /// May return errors for:
504 /// - Invalid UTF-8 sequences
505 /// - Text exceeding maximum length
506 /// - Unknown tokens that cannot be handled
507 fn encode(&self, text: &str) -> Result<TokenizedInput>;
508
509 /// Encodes a pair of texts for sequence-pair tasks.
510 ///
511 /// This method is used for tasks that require two input sequences,
512 /// such as question answering, textual entailment, or sequence classification.
513 ///
514 /// # Arguments
515 ///
516 /// * `text` - The first text sequence
517 /// * `text2` - The second text sequence
518 ///
519 /// # Returns
520 ///
521 /// Returns a `TokenizedInput` with both sequences encoded and separated
522 /// by appropriate special tokens (e.g., `[SEP]` for BERT).
523 ///
524 /// # Errors
525 ///
526 /// May return errors for the same reasons as `encode()`.
527 fn encode_pair(&self, text: &str, text2: &str) -> Result<TokenizedInput>;
528
529 /// Decodes token IDs back into text.
530 ///
531 /// # Arguments
532 ///
533 /// * `ids` - The token IDs to decode
534 ///
535 /// # Returns
536 ///
537 /// Returns the decoded text string. Special tokens may be included
538 /// or excluded depending on the implementation.
539 ///
540 /// # Errors
541 ///
542 /// May return errors for:
543 /// - Invalid token IDs
544 /// - Decoding errors
545 fn decode(&self, ids: &[u32]) -> Result<String>;
546
547 /// Returns the size of the tokenizer's vocabulary.
548 ///
549 /// # Returns
550 ///
551 /// The total number of tokens in the vocabulary.
552 fn vocab_size(&self) -> usize;
553
554 /// Returns a copy of the vocabulary as a mapping from tokens to IDs.
555 ///
556 /// # Returns
557 ///
558 /// A HashMap containing the vocabulary mapping.
559 fn get_vocab(&self) -> std::collections::HashMap<String, u32>;
560
561 /// Converts a token string to its corresponding ID.
562 ///
563 /// # Arguments
564 ///
565 /// * `token` - The token string to convert
566 ///
567 /// # Returns
568 ///
569 /// The token ID if the token exists in the vocabulary, None otherwise.
570 fn token_to_id(&self, token: &str) -> Option<u32>;
571
572 /// Converts a token ID to its corresponding token string.
573 ///
574 /// # Arguments
575 ///
576 /// * `id` - The token ID to convert
577 ///
578 /// # Returns
579 ///
580 /// The token string if the ID exists in the vocabulary, None otherwise.
581 fn id_to_token(&self, id: u32) -> Option<String>;
582}
583
584/// Represents tokenized input ready for model consumption.
585///
586/// `TokenizedInput` contains all the necessary components for feeding
587/// text data into a transformer model after tokenization.
588///
589/// # Fields
590///
591/// * `input_ids` - The token IDs representing the input text
592/// * `attention_mask` - Binary mask (0 or 1) indicating which tokens are real vs padding
593/// * `token_type_ids` - Optional segment IDs for models that use them (e.g., BERT)
594///
595/// # Example
596///
597/// ```no_run
598/// use trustformers_core::traits::TokenizedInput;
599///
600/// let input = TokenizedInput::with_token_type_ids(
601/// vec![101, 2023, 2003, 1037, 3231, 102], // [CLS] this is a test [SEP]
602/// vec![1, 1, 1, 1, 1, 1], // All tokens are real (not padding)
603/// Some(vec![0, 0, 0, 0, 0, 0]), // All tokens from first segment
604/// );
605/// ```
606#[derive(Debug, Clone, Default)]
607pub struct TokenizedInput {
608 /// Token IDs representing the encoded text.
609 /// These correspond to entries in the tokenizer's vocabulary.
610 pub input_ids: Vec<u32>,
611
612 /// Binary attention mask indicating real tokens (1) vs padding tokens (0).
613 /// This prevents the model from attending to padding tokens.
614 pub attention_mask: Vec<u8>,
615
616 /// Optional token type IDs for distinguishing between different segments.
617 /// Used by models like BERT for tasks involving multiple sequences.
618 /// Typically 0 for the first sequence and 1 for the second sequence.
619 pub token_type_ids: Option<Vec<u32>>,
620
621 /// Optional special tokens mask indicating special tokens (1) vs regular tokens (0).
622 /// Used to identify tokens like `[CLS]`, `[SEP]`, `[PAD]` etc.
623 pub special_tokens_mask: Option<Vec<u8>>,
624
625 /// Optional offset mapping showing character positions of tokens in original text.
626 /// Each tuple contains (start_pos, end_pos) character offsets.
627 pub offset_mapping: Option<Vec<(usize, usize)>>,
628
629 /// Optional overflowing tokens when text exceeds max length.
630 /// Contains tokens that were truncated from the input.
631 pub overflowing_tokens: Option<Vec<u32>>,
632}
633
634impl TokenizedInput {
635 /// Create a new TokenizedInput with minimal required fields
636 pub fn new(input_ids: Vec<u32>, attention_mask: Vec<u8>) -> Self {
637 Self {
638 input_ids,
639 attention_mask,
640 token_type_ids: None,
641 special_tokens_mask: None,
642 offset_mapping: None,
643 overflowing_tokens: None,
644 }
645 }
646
647 /// Create a new TokenizedInput with token type IDs
648 pub fn with_token_type_ids(
649 input_ids: Vec<u32>,
650 attention_mask: Vec<u8>,
651 token_type_ids: Option<Vec<u32>>,
652 ) -> Self {
653 Self {
654 input_ids,
655 attention_mask,
656 token_type_ids,
657 special_tokens_mask: None,
658 offset_mapping: None,
659 overflowing_tokens: None,
660 }
661 }
662}
663
664/// Parameter optimization algorithms for training neural networks.
665///
666/// The `Optimizer` trait defines the interface for gradient-based optimization
667/// algorithms such as SGD, Adam, AdamW, etc. Optimizers update model parameters
668/// based on computed gradients to minimize the loss function.
669///
670/// # Thread Safety
671///
672/// Optimizers must be `Send + Sync` to support distributed training.
673///
674/// # Example
675///
676/// ```no_run
677/// use trustformers_core::traits::Optimizer;
678/// use trustformers_core::tensor::Tensor;
679/// use trustformers_core::errors::Result;
680///
681/// struct SGD {
682/// learning_rate: f32,
683/// momentum: f32,
684/// velocity: std::collections::HashMap<String, Tensor>,
685/// }
686///
687/// impl Optimizer for SGD {
688/// fn update(&mut self, _parameter: &mut Tensor, _grad: &Tensor) -> Result<()> {
689/// // SGD with momentum: v = momentum * v - lr * grad
690/// // parameter += v
691/// Ok(())
692/// }
693///
694/// fn zero_grad(&mut self) {
695/// // Clear accumulated gradients
696/// }
697///
698/// fn step(&mut self) {
699/// // Apply updates to all parameters
700/// }
701///
702/// fn get_lr(&self) -> f32 {
703/// self.learning_rate
704/// }
705///
706/// fn set_lr(&mut self, lr: f32) {
707/// self.learning_rate = lr;
708/// }
709/// }
710/// ```
711pub trait Optimizer: Send + Sync {
712 /// Updates a parameter based on its gradient.
713 ///
714 /// # Arguments
715 ///
716 /// * `parameter` - The parameter tensor to update
717 /// * `grad` - The gradient tensor for this parameter
718 ///
719 /// # Returns
720 ///
721 /// Returns `Ok(())` on successful update, or an error if the update fails.
722 ///
723 /// # Errors
724 ///
725 /// May return errors for:
726 /// - Mismatched tensor shapes
727 /// - Numerical errors (e.g., NaN or Inf values)
728 /// - Memory allocation failures
729 fn update(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()>;
730
731 /// Clears all accumulated gradients.
732 ///
733 /// This should be called before each backward pass to ensure
734 /// gradients don't accumulate across batches (unless gradient
735 /// accumulation is intentionally being used).
736 fn zero_grad(&mut self);
737
738 /// Performs a single optimization step.
739 ///
740 /// This method applies all pending parameter updates. It should be
741 /// called after gradients have been computed for all parameters.
742 fn step(&mut self);
743
744 /// Gets the current learning rate.
745 ///
746 /// # Returns
747 ///
748 /// The current learning rate value.
749 fn get_lr(&self) -> f32;
750
751 /// Sets a new learning rate.
752 ///
753 /// # Arguments
754 ///
755 /// * `lr` - The new learning rate value
756 ///
757 /// # Note
758 ///
759 /// This is useful for implementing learning rate schedules.
760 fn set_lr(&mut self, lr: f32);
761
762 /// Accumulates gradients for gradient accumulation.
763 ///
764 /// This method is used when training with gradient accumulation,
765 /// where gradients from multiple batches are accumulated before
766 /// performing an update step.
767 ///
768 /// # Arguments
769 ///
770 /// * `parameter` - The parameter tensor
771 /// * `grad` - The gradient to accumulate
772 ///
773 /// # Returns
774 ///
775 /// Returns `Ok(())` on success, or an error if accumulation fails.
776 ///
777 /// # Default Implementation
778 ///
779 /// The default implementation simply calls `update()`. Override this
780 /// method for optimizers that need special gradient accumulation logic.
781 fn accumulate_grad(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()> {
782 // Default implementation: just store the gradient for later use
783 self.update(parameter, grad)
784 }
785
786 /// Applies accumulated gradients after gradient accumulation.
787 ///
788 /// This method should be called after accumulating gradients from
789 /// multiple batches to apply the averaged update.
790 ///
791 /// # Arguments
792 ///
793 /// * `accumulation_steps` - The number of accumulation steps performed
794 ///
795 /// # Returns
796 ///
797 /// Returns `Ok(())` on success, or an error if application fails.
798 ///
799 /// # Default Implementation
800 ///
801 /// The default implementation is a no-op. Override this method for
802 /// optimizers that implement gradient accumulation.
803 fn apply_accumulated_grads(&mut self, accumulation_steps: usize) -> Result<()> {
804 // Default implementation: no-op, override if needed
805 let _ = accumulation_steps;
806 Ok(())
807 }
808}
809
810/// Weight initialization strategies for neural network parameters.
811///
812/// The `ParameterInit` trait provides various initialization methods that help
813/// ensure proper gradient flow and training stability. Different initialization
814/// strategies are optimal for different activation functions and architectures.
815///
816/// # Example
817///
818/// ```no_run
819/// use trustformers_core::traits::ParameterInit;
820///
821/// // Example type implementing ParameterInit
822/// struct WeightMatrix {
823/// data: Vec<f32>,
824/// shape: [usize; 2],
825/// }
826///
827/// impl ParameterInit for WeightMatrix {
828/// fn normal(&mut self, mean: f32, std: f32) {
829/// // fill data with normal distribution values
830/// }
831/// fn uniform(&mut self, min: f32, max: f32) {
832/// // fill data with uniform distribution values
833/// }
834/// fn xavier_uniform(&mut self) {
835/// // Xavier/Glorot initialization
836/// }
837/// fn xavier_normal(&mut self) {}
838/// fn kaiming_uniform(&mut self, _mode: &str, _nonlinearity: &str) {}
839/// fn kaiming_normal(&mut self, _mode: &str, _nonlinearity: &str) {}
840/// }
841///
842/// let mut weight = WeightMatrix { data: vec![0.0; 768 * 768], shape: [768, 768] };
843///
844/// // Initialize with Xavier/Glorot uniform for tanh activations
845/// weight.xavier_uniform();
846///
847/// // Or use Kaiming/He initialization for ReLU activations
848/// weight.kaiming_normal("fan_in", "relu");
849/// ```
850pub trait ParameterInit {
851 /// Initializes the tensor with values from a normal distribution.
852 ///
853 /// # Arguments
854 ///
855 /// * `mean` - The mean of the normal distribution
856 /// * `std` - The standard deviation of the normal distribution
857 ///
858 /// # Example
859 ///
860 /// ```no_run
861 /// # use trustformers_core::traits::ParameterInit;
862 /// # struct Weights { data: Vec<f32> }
863 /// # impl ParameterInit for Weights {
864 /// # fn normal(&mut self, mean: f32, std: f32) {}
865 /// # fn uniform(&mut self, min: f32, max: f32) {}
866 /// # fn xavier_uniform(&mut self) {}
867 /// # fn xavier_normal(&mut self) {}
868 /// # fn kaiming_uniform(&mut self, _: &str, _: &str) {}
869 /// # fn kaiming_normal(&mut self, _: &str, _: &str) {}
870 /// # }
871 /// let mut tensor = Weights { data: vec![0.0; 10000] };
872 /// tensor.normal(0.0, 0.02); // Common for transformer embeddings
873 /// ```
874 fn normal(&mut self, mean: f32, std: f32);
875
876 /// Initializes the tensor with values from a uniform distribution.
877 ///
878 /// # Arguments
879 ///
880 /// * `min` - The minimum value (inclusive)
881 /// * `max` - The maximum value (exclusive)
882 ///
883 /// # Example
884 ///
885 /// ```no_run
886 /// # use trustformers_core::traits::ParameterInit;
887 /// # struct Weights { data: Vec<f32> }
888 /// # impl ParameterInit for Weights {
889 /// # fn normal(&mut self, mean: f32, std: f32) {}
890 /// # fn uniform(&mut self, min: f32, max: f32) {}
891 /// # fn xavier_uniform(&mut self) {}
892 /// # fn xavier_normal(&mut self) {}
893 /// # fn kaiming_uniform(&mut self, _: &str, _: &str) {}
894 /// # fn kaiming_normal(&mut self, _: &str, _: &str) {}
895 /// # }
896 /// let mut tensor = Weights { data: vec![0.0; 10000] };
897 /// tensor.uniform(-0.1, 0.1);
898 /// ```
899 fn uniform(&mut self, min: f32, max: f32);
900
901 /// Xavier/Glorot uniform initialization.
902 ///
903 /// Initializes weights to maintain variance across layers, optimal for
904 /// tanh and sigmoid activations. The range is [-x, x] where
905 /// x = sqrt(6 / (fan_in + fan_out)).
906 ///
907 /// # References
908 ///
909 /// Glorot & Bengio (2010): "Understanding the difficulty of training
910 /// deep feedforward neural networks"
911 fn xavier_uniform(&mut self);
912
913 /// Xavier/Glorot normal initialization.
914 ///
915 /// Similar to `xavier_uniform` but uses a normal distribution with
916 /// std = sqrt(2 / (fan_in + fan_out)).
917 fn xavier_normal(&mut self);
918
919 /// Kaiming/He uniform initialization.
920 ///
921 /// Designed for ReLU and similar activations. Maintains variance when
922 /// half of the neurons are zeroed out by ReLU.
923 ///
924 /// # Arguments
925 ///
926 /// * `mode` - Either "fan_in" or "fan_out", determines which dimension to use
927 /// * `nonlinearity` - The activation function ("relu", "leaky_relu", "linear")
928 ///
929 /// # Example
930 ///
931 /// ```no_run
932 /// # use trustformers_core::traits::ParameterInit;
933 /// # struct Weights { data: Vec<f32> }
934 /// # impl ParameterInit for Weights {
935 /// # fn normal(&mut self, mean: f32, std: f32) {}
936 /// # fn uniform(&mut self, min: f32, max: f32) {}
937 /// # fn xavier_uniform(&mut self) {}
938 /// # fn xavier_normal(&mut self) {}
939 /// # fn kaiming_uniform(&mut self, _: &str, _: &str) {}
940 /// # fn kaiming_normal(&mut self, _: &str, _: &str) {}
941 /// # }
942 /// let mut conv_weight = Weights { data: vec![0.0; 64 * 32 * 3 * 3] };
943 /// conv_weight.kaiming_uniform("fan_in", "relu");
944 /// ```
945 ///
946 /// # References
947 ///
948 /// He et al. (2015): "Delving Deep into Rectifiers: Surpassing
949 /// Human-Level Performance on ImageNet Classification"
950 fn kaiming_uniform(&mut self, mode: &str, nonlinearity: &str);
951
952 /// Kaiming/He normal initialization.
953 ///
954 /// Similar to `kaiming_uniform` but uses a normal distribution.
955 /// Generally preferred over uniform for deeper networks.
956 ///
957 /// # Arguments
958 ///
959 /// * `mode` - Either "fan_in" or "fan_out"
960 /// * `nonlinearity` - The activation function ("relu", "leaky_relu", "linear")
961 fn kaiming_normal(&mut self, mode: &str, nonlinearity: &str);
962}