Struct TransformerBlock 

pub struct TransformerBlock { /* private fields */ }

Transformer block for TRM

Each block consists of:

• Self-attention (optional, not used in MLP mode)
• Feed-forward network (SwiGLU)
• RMS normalization with residual connections (post-norm)

Implementations

impl TransformerBlock

pub fn new(config: TRMConfig, vb: VarBuilder<'_>) -> Result<Self>

Create new transformer block

pub fn forward( &self, hidden_states: &Tensor, cos_sin: Option<(&Tensor, &Tensor)>, ) -> Result<Tensor>

Forward pass

Arguments

• hidden_states - Input tensor [batch, seq_len, hidden_size]
• cos_sin - Optional RoPE embeddings

Returns

Output tensor [batch, seq_len, hidden_size]