oa::ml::nn::TransformerBlock

One pre-normalized causal Transformer block.

StRust public API

Declaration

pub struct TransformerBlock { /* private fields */ }
Public Rust declaration.

Public Methods

pub fn forward(&self, input: &Matrix) -> Result<Matrix>
pub fn forward_masked( &self, input: &Matrix, additive_mask: &Matrix, ) -> Result<Matrix>
pub fn forward_conditioned( &self, input: &Matrix, condition: &Matrix, ) -> Result<Matrix>
pub fn forward_conditioned_masked( &self, input: &Matrix, condition: &Matrix, additive_mask: &Matrix, ) -> Result<Matrix>
pub const fn model_width(&self) -> usize
pub const fn hidden_width(&self) -> usize
pub fn sequence_length(&self) -> usize
pub const fn num_heads(&self) -> usize
pub fn is_moe(&self) -> bool
pub fn moe(&self) -> Option<&Moe>
pub fn is_adaptively_conditioned(&self) -> bool
pub const fn condition_dim(&self) -> Option<usize>
pub fn set_attention_mode(&self, mode: AttentionMode)
pub fn attention_mode(&self) -> AttentionMode
pub fn set_sequence_length(&self, sequence_length: usize) -> Result<()>
fn forward(&self, input: &Matrix) -> Result<Matrix>
fn registry(&self) -> &ModuleRegistry
fn parameters(&self) -> Vec<Parameter>
fn named_parameters(&self) -> Vec<NamedParameter>
fn buffers(&self) -> Vec<Matrix>
fn named_buffers(&self) -> Vec<NamedBuffer>
fn all_parameters(&self) -> Result<Vec<Parameter>>
fn all_named_parameters(&self) -> Result<Vec<NamedParameter>>
fn all_named_buffers(&self) -> Result<Vec<NamedBuffer>>
fn all_named_state_u32(&self) -> Result<Vec<NamedStateU32>>
fn num_parameters(&self) -> Result<usize>
fn train(&self, training: bool)
fn eval(&self)
fn is_training(&self) -> bool
fn scoped_eval(&self) -> ScopedEval<'_>

Associated Functions

pub fn with_seed( engine: &Engine, model_width: usize, hidden_width: usize, sequence_length: usize, num_heads: usize, epsilon: f32, seed: u64, ) -> Result<Self>
pub fn with_seed_conditioned( engine: &Engine, model_width: usize, hidden_width: usize, sequence_length: usize, num_heads: usize, condition_dim: usize, epsilon: f32, seed: u64, ) -> Result<Self>
pub fn with_seed_moe( engine: &Engine, model_width: usize, expert_hidden_width: usize, sequence_length: usize, num_heads: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, seed: u64, ) -> Result<Self>
pub fn with_seed_moe_conditioned( engine: &Engine, model_width: usize, expert_hidden_width: usize, sequence_length: usize, num_heads: usize, num_experts: usize, experts_per_token: usize, condition_dim: usize, epsilon: f32, seed: u64, ) -> Result<Self>

Public Method Documentation

pub fn forward( &self, input: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
pub fn forward_masked( &self, input: &Matrix, additive_mask: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
pub fn forward_conditioned( &self, input: &Matrix, condition: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
pub fn forward_conditioned_masked( &self, input: &Matrix, condition: &Matrix, additive_mask: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
pub const fn model_width( &self ) -> usize
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
pub const fn hidden_width( &self ) -> usize
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
pub fn sequence_length( &self ) -> usize
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
pub const fn num_heads( &self ) -> usize
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
pub fn is_moe( &self ) -> bool
No public source comment is attached to this declaration.
pub fn moe( &self ) -> Option<&Moe>
No public source comment is attached to this declaration.
pub fn is_adaptively_conditioned( &self ) -> bool
No public source comment is attached to this declaration.
pub const fn condition_dim( &self ) -> Option<usize>
No public source comment is attached to this declaration.
pub fn set_attention_mode( &self, mode: AttentionMode )
No public source comment is attached to this declaration.
pub fn attention_mode( &self ) -> AttentionMode
No public source comment is attached to this declaration.
pub fn set_sequence_length(&self, sequence_length: usize) -> Result<()>
No public source comment is attached to this declaration.
fn forward( &self, input: &Matrix ) -> Result<Matrix>
No public source comment is attached to this declaration.
fn registry( &self ) -> &ModuleRegistry
No public source comment is attached to this declaration.
fn parameters( &self ) -> Vec<Parameter>
No public source comment is attached to this declaration.
fn named_parameters( &self ) -> Vec<NamedParameter>
No public source comment is attached to this declaration.
fn buffers( &self ) -> Vec<Matrix>
No public source comment is attached to this declaration.
fn named_buffers( &self ) -> Vec<NamedBuffer>
No public source comment is attached to this declaration.
fn all_parameters( &self ) -> Result<Vec<Parameter>>
No public source comment is attached to this declaration.
fn all_named_parameters( &self ) -> Result<Vec<NamedParameter>>
No public source comment is attached to this declaration.
fn all_named_buffers( &self ) -> Result<Vec<NamedBuffer>>
No public source comment is attached to this declaration.
fn all_named_state_u32( &self ) -> Result<Vec<NamedStateU32>>
No public source comment is attached to this declaration.
fn num_parameters( &self ) -> Result<usize>
No public source comment is attached to this declaration.
fn train( &self, training: bool )
No public source comment is attached to this declaration.
fn eval( &self )
No public source comment is attached to this declaration.
fn is_training( &self ) -> bool
No public source comment is attached to this declaration.
fn scoped_eval( &self ) -> ScopedEval<'_>
No public source comment is attached to this declaration.
pub fn with_seed( engine: &Engine, model_width: usize, hidden_width: usize, sequence_length: usize, num_heads: usize, epsilon: f32, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.
pub fn with_seed_conditioned( engine: &Engine, model_width: usize, hidden_width: usize, sequence_length: usize, num_heads: usize, condition_dim: usize, epsilon: f32, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.
pub fn with_seed_moe( engine: &Engine, model_width: usize, expert_hidden_width: usize, sequence_length: usize, num_heads: usize, num_experts: usize, experts_per_token: usize, epsilon: f32, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.
pub fn with_seed_moe_conditioned( engine: &Engine, model_width: usize, expert_hidden_width: usize, sequence_length: usize, num_heads: usize, num_experts: usize, experts_per_token: usize, condition_dim: usize, epsilon: f32, seed: u64 ) -> Result<Self>
No public source comment is attached to this declaration.