Struct TransformerBlock 

pub struct TransformerBlock<B>
where
    B: Backend,
{ /* private fields */ }

A single Transformer block (pre-norm style).

Contains:

• Self-attention with multi-head attention
• Feed-forward network (two linear layers with GELU)
• Two layer normalizations
• Residual connections around both sub-layers

Implementations

impl<B> TransformerBlock<B>

where B: Backend,

pub fn new( d_model: usize, num_heads: usize, d_ff: usize, causal: bool, dtype: DType, device: &<B as Backend>::Device, ) -> Result<TransformerBlock<B>, Error>

Create a new TransformerBlock.

Arguments

• d_model: model dimension (embedding size)
• num_heads: number of attention heads
• d_ff: feed-forward inner dimension (typically 4 * d_model)
• causal: whether to use causal (autoregressive) attention mask
• dtype: data type
• device: compute device

pub fn d_model(&self) -> usize

Trait Implementations

impl<B> Module<B> for TransformerBlock<B>

where B: Backend,

fn forward(&self, x: &Tensor<B>) -> Result<Tensor<B>, Error>

Forward pass (pre-norm): x = x + attention(layernorm(x)) x = x + ffn(layernorm(x))

fn parameters(&self) -> Vec<Tensor<B>>

Return all trainable parameters of this module. The optimizer uses these to update weights during training.

fn named_parameters(&self) -> Vec<(String, Tensor<B>)>

Return all trainable parameters with human-readable names.

fn set_training(&self, _training: bool)

Set training or evaluation mode.

fn is_training(&self) -> bool

Whether the module is in training mode (default: true).

fn train(&self)

Convenience: set training mode.

fn eval(&self)

Convenience: set evaluation mode.

fn num_parameters(&self) -> usize

Total number of scalar parameters in this module.

fn trainable_params_count(&self) -> usize

Number of trainable (variable) parameters.

fn frozen_parameters(&self) -> Vec<Tensor<B>>

Freeze all parameters: returns new parameter tensors with is_variable = false, preventing gradient accumulation.

fn state_dict(&self) -> Vec<(String, Tensor<B>)>

Returns a state_dict-style map of parameter name → tensor.