Struct ReasoningModule 

pub struct ReasoningModule { /* private fields */ }

Reasoning module: stack of transformer blocks with input injection

This is the L-level or H-level reasoning component.

Implementations

impl ReasoningModule

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

Create new reasoning module

Arguments

• num_layers - Number of transformer blocks
• config - Model configuration
• vb - VarBuilder for parameter initialization

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

Forward pass with input injection

Arguments

• hidden_states - Current state
• input_injection - Input to inject (added to hidden_states)
• cos_sin - Optional RoPE embeddings

Returns

Updated state after processing through all layers