Struct Transformer

pub struct Transformer<F: Float + Debug + Send + Sync> { /* private fields */ }

Complete transformer model with encoder and decoder

This implements the full transformer architecture from “Attention Is All You Need”, combining encoder and decoder stacks with positional encoding.

Implementations

impl<F: Float + Debug + ScalarOperand + Send + Sync + 'static> Transformer<F>

pub fn new<R: Rng>(config: TransformerConfig, rng: &mut R) -> Result<Self>

Create a new transformer model

Arguments

• config - Transformer configuration
• rng - Random number generator for weight initialization

Returns

• A new transformer model

Examples found in repository

examples/transformer_example.rs (line 42)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Transformer Model Example");
    println!("========================");

    // Create a seeded RNG for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a small transformer configuration for demonstration
    let config = TransformerConfig {
        d_model: 64,                                           // Embedding dimension
        n_encoder_layers: 2,                                   // Number of encoder layers
        n_decoder_layers: 2,                                   // Number of decoder layers
        n_heads: 4,                                            // Number of attention heads
        d_ff: 128,       // Feed-forward network hidden dimension
        max_seq_len: 50, // Maximum sequence length
        dropout: 0.1,    // Dropout rate
        pos_encoding_type: PositionalEncodingType::Sinusoidal, // Positional encoding type
        epsilon: 1e-5,   // Small constant for layer normalization
    };

    println!("Creating transformer model with config:");
    println!("  - d_model: {}", config.d_model);
    println!("  - n_encoder_layers: {}", config.n_encoder_layers);
    println!("  - n_decoder_layers: {}", config.n_decoder_layers);
    println!("  - n_heads: {}", config.n_heads);
    println!("  - d_ff: {}", config.d_ff);
    println!("  - max_seq_len: {}", config.max_seq_len);

    // Create the transformer model
    let transformer = Transformer::<f64>::new(config, &mut rng)?;

    // Create sample inputs
    // In a real application, these would be token embeddings
    let batch_size = 2;
    let src_seq_len = 10;
    let tgt_seq_len = 8;
    let d_model = 64;

    println!("\nSample dimensions:");
    println!("  - Batch size: {}", batch_size);
    println!("  - Source sequence length: {}", src_seq_len);
    println!("  - Target sequence length: {}", tgt_seq_len);

    // Create source and target sequence embeddings
    // In practice, these would come from embedding layers or tokenizers
    let src_embeddings = Array3::<f64>::from_elem((batch_size, src_seq_len, d_model), 0.1);
    let tgt_embeddings = Array3::<f64>::from_elem((batch_size, tgt_seq_len, d_model), 0.1);

    // Convert to dyn format once and reuse
    let src_embeddings_dyn = src_embeddings.clone().into_dyn();
    let tgt_embeddings_dyn = tgt_embeddings.clone().into_dyn();

    println!("\nRunning encoder-only inference...");
    // Run encoder-only inference (useful for tasks like classification)
    let encoder_output = transformer.forward(&src_embeddings_dyn)?;
    println!("Encoder output shape: {:?}", encoder_output.shape());

    println!("\nRunning full transformer inference (training mode)...");
    // Run full transformer training (teacher forcing)
    let output_train = transformer.forward_train(&src_embeddings_dyn, &tgt_embeddings_dyn)?;
    println!("Training output shape: {:?}", output_train.shape());

    println!("\nRunning autoregressive inference (one step)...");
    // Simulate autoregressive generation (one step)
    // In practice, we would use a loop to generate tokens one by one
    let first_token = Array3::<f64>::from_elem((batch_size, 1, d_model), 0.1);
    let first_token_dyn = first_token.clone().into_dyn();
    let output_inference = transformer.forward_inference(&src_embeddings_dyn, &first_token_dyn)?;
    println!("Inference output shape: {:?}", output_inference.shape());

    println!("\nExample completed successfully");
    Ok(())
}

pub fn forward_train( &self, src: &Array<F, IxDyn>, tgt: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Forward pass with encoder and decoder

Arguments

• src - Source sequences [batch, src_len, d_model]
• tgt - Target sequences [batch, tgt_len, d_model]

Returns

• Output tensor [batch, tgt_len, d_model]

Examples found in repository

examples/transformer_example.rs (line 72)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Transformer Model Example");
    println!("========================");

    // Create a seeded RNG for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a small transformer configuration for demonstration
    let config = TransformerConfig {
        d_model: 64,                                           // Embedding dimension
        n_encoder_layers: 2,                                   // Number of encoder layers
        n_decoder_layers: 2,                                   // Number of decoder layers
        n_heads: 4,                                            // Number of attention heads
        d_ff: 128,       // Feed-forward network hidden dimension
        max_seq_len: 50, // Maximum sequence length
        dropout: 0.1,    // Dropout rate
        pos_encoding_type: PositionalEncodingType::Sinusoidal, // Positional encoding type
        epsilon: 1e-5,   // Small constant for layer normalization
    };

    println!("Creating transformer model with config:");
    println!("  - d_model: {}", config.d_model);
    println!("  - n_encoder_layers: {}", config.n_encoder_layers);
    println!("  - n_decoder_layers: {}", config.n_decoder_layers);
    println!("  - n_heads: {}", config.n_heads);
    println!("  - d_ff: {}", config.d_ff);
    println!("  - max_seq_len: {}", config.max_seq_len);

    // Create the transformer model
    let transformer = Transformer::<f64>::new(config, &mut rng)?;

    // Create sample inputs
    // In a real application, these would be token embeddings
    let batch_size = 2;
    let src_seq_len = 10;
    let tgt_seq_len = 8;
    let d_model = 64;

    println!("\nSample dimensions:");
    println!("  - Batch size: {}", batch_size);
    println!("  - Source sequence length: {}", src_seq_len);
    println!("  - Target sequence length: {}", tgt_seq_len);

    // Create source and target sequence embeddings
    // In practice, these would come from embedding layers or tokenizers
    let src_embeddings = Array3::<f64>::from_elem((batch_size, src_seq_len, d_model), 0.1);
    let tgt_embeddings = Array3::<f64>::from_elem((batch_size, tgt_seq_len, d_model), 0.1);

    // Convert to dyn format once and reuse
    let src_embeddings_dyn = src_embeddings.clone().into_dyn();
    let tgt_embeddings_dyn = tgt_embeddings.clone().into_dyn();

    println!("\nRunning encoder-only inference...");
    // Run encoder-only inference (useful for tasks like classification)
    let encoder_output = transformer.forward(&src_embeddings_dyn)?;
    println!("Encoder output shape: {:?}", encoder_output.shape());

    println!("\nRunning full transformer inference (training mode)...");
    // Run full transformer training (teacher forcing)
    let output_train = transformer.forward_train(&src_embeddings_dyn, &tgt_embeddings_dyn)?;
    println!("Training output shape: {:?}", output_train.shape());

    println!("\nRunning autoregressive inference (one step)...");
    // Simulate autoregressive generation (one step)
    // In practice, we would use a loop to generate tokens one by one
    let first_token = Array3::<f64>::from_elem((batch_size, 1, d_model), 0.1);
    let first_token_dyn = first_token.clone().into_dyn();
    let output_inference = transformer.forward_inference(&src_embeddings_dyn, &first_token_dyn)?;
    println!("Inference output shape: {:?}", output_inference.shape());

    println!("\nExample completed successfully");
    Ok(())
}

pub fn forward_inference( &self, src: &Array<F, IxDyn>, tgt: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Forward pass for inference (without target)

Arguments

• src - Source sequences [batch, src_len, d_model]
• tgt - Target sequences so far [batch, tgt_len, d_model]

Returns

• Output tensor [batch, tgt_len, d_model]

Examples found in repository

examples/transformer_example.rs (line 80)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Transformer Model Example");
    println!("========================");

    // Create a seeded RNG for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a small transformer configuration for demonstration
    let config = TransformerConfig {
        d_model: 64,                                           // Embedding dimension
        n_encoder_layers: 2,                                   // Number of encoder layers
        n_decoder_layers: 2,                                   // Number of decoder layers
        n_heads: 4,                                            // Number of attention heads
        d_ff: 128,       // Feed-forward network hidden dimension
        max_seq_len: 50, // Maximum sequence length
        dropout: 0.1,    // Dropout rate
        pos_encoding_type: PositionalEncodingType::Sinusoidal, // Positional encoding type
        epsilon: 1e-5,   // Small constant for layer normalization
    };

    println!("Creating transformer model with config:");
    println!("  - d_model: {}", config.d_model);
    println!("  - n_encoder_layers: {}", config.n_encoder_layers);
    println!("  - n_decoder_layers: {}", config.n_decoder_layers);
    println!("  - n_heads: {}", config.n_heads);
    println!("  - d_ff: {}", config.d_ff);
    println!("  - max_seq_len: {}", config.max_seq_len);

    // Create the transformer model
    let transformer = Transformer::<f64>::new(config, &mut rng)?;

    // Create sample inputs
    // In a real application, these would be token embeddings
    let batch_size = 2;
    let src_seq_len = 10;
    let tgt_seq_len = 8;
    let d_model = 64;

    println!("\nSample dimensions:");
    println!("  - Batch size: {}", batch_size);
    println!("  - Source sequence length: {}", src_seq_len);
    println!("  - Target sequence length: {}", tgt_seq_len);

    // Create source and target sequence embeddings
    // In practice, these would come from embedding layers or tokenizers
    let src_embeddings = Array3::<f64>::from_elem((batch_size, src_seq_len, d_model), 0.1);
    let tgt_embeddings = Array3::<f64>::from_elem((batch_size, tgt_seq_len, d_model), 0.1);

    // Convert to dyn format once and reuse
    let src_embeddings_dyn = src_embeddings.clone().into_dyn();
    let tgt_embeddings_dyn = tgt_embeddings.clone().into_dyn();

    println!("\nRunning encoder-only inference...");
    // Run encoder-only inference (useful for tasks like classification)
    let encoder_output = transformer.forward(&src_embeddings_dyn)?;
    println!("Encoder output shape: {:?}", encoder_output.shape());

    println!("\nRunning full transformer inference (training mode)...");
    // Run full transformer training (teacher forcing)
    let output_train = transformer.forward_train(&src_embeddings_dyn, &tgt_embeddings_dyn)?;
    println!("Training output shape: {:?}", output_train.shape());

    println!("\nRunning autoregressive inference (one step)...");
    // Simulate autoregressive generation (one step)
    // In practice, we would use a loop to generate tokens one by one
    let first_token = Array3::<f64>::from_elem((batch_size, 1, d_model), 0.1);
    let first_token_dyn = first_token.clone().into_dyn();
    let output_inference = transformer.forward_inference(&src_embeddings_dyn, &first_token_dyn)?;
    println!("Inference output shape: {:?}", output_inference.shape());

    println!("\nExample completed successfully");
    Ok(())
}

Trait Implementations

impl<F: Float + Debug + ScalarOperand + Send + Sync + 'static> Layer<F> for Transformer<F>

fn as_any(&self) -> &dyn Any

Get the layer as a dyn Any for downcasting

fn as_any_mut(&mut self) -> &mut dyn Any

Get the layer as a mutable dyn Any for downcasting

fn forward(&self, input: &Array<F, IxDyn>) -> Result<Array<F, IxDyn>>

Forward pass of the layer

fn backward( &self, input: &Array<F, IxDyn>, _grad_output: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Backward pass of the layer to compute gradients

fn update(&mut self, learning_rate: F) -> Result<()>

Update the layer parameters with the given gradients

fn params(&self) -> Vec<Array<F, IxDyn>>

Get the parameters of the layer

fn gradients(&self) -> Vec<Array<F, IxDyn>>

Get the gradients of the layer parameters

fn set_gradients(&mut self, _gradients: &[Array<F, IxDyn>]) -> Result<()>

Set the gradients of the layer parameters

fn set_params(&mut self, _params: &[Array<F, IxDyn>]) -> Result<()>

Set the parameters of the layer

fn set_training(&mut self, _training: bool)

Set the layer to training mode (true) or evaluation mode (false)

fn is_training(&self) -> bool

Get the current training mode

fn layer_type(&self) -> &str

Get the type of the layer (e.g., “Dense”, “Conv2D”)

fn parameter_count(&self) -> usize

Get the number of trainable parameters in this layer

fn layer_description(&self) -> String

Get a detailed description of this layer