entrenar/transformer/
mod.rs1mod attention;
12mod block;
13mod config;
14pub(crate) mod cuda_block;
15mod embedding;
16mod encoder;
17mod encoder_block;
18mod feedforward;
19pub mod init;
20mod model;
21mod norm;
22pub(crate) mod weights;
23#[cfg(feature = "gpu")]
24pub mod wgpu_block;
25
26pub use attention::{LoRAProjection, MultiHeadAttention, MultiHeadAttentionWithLoRA};
27pub use config::{ModelArchitecture, TransformerConfig};
28#[cfg(feature = "cuda")]
29pub use cuda_block::{
30 BlockWeights, CudaBlock, CudaGradWorkspace, CudaNf4TransformerBlock, CudaTransformerBlock,
31 GpuBlockOptimizerState,
32};
33#[cfg(not(feature = "cuda"))]
34pub use cuda_block::{CudaBlock, CudaTransformerBlock};
35#[cfg(feature = "cuda")]
36pub(crate) use cuda_block::{CudaBlockScratch, CudaLoraGradWorkspace, GpuLoraOptimizerState};
37pub use embedding::LearnedPositionEmbedding;
38pub use encoder::EncoderModel;
39pub use encoder_block::EncoderBlock;
40pub use feedforward::EncoderFeedForward;
41pub use model::Transformer;
42pub use norm::LayerNorm;
43pub use weights::{load_safetensors_weights, validate_weights, Architecture};
44#[cfg(feature = "gpu")]
45pub use wgpu_block::WgpuForwardPass;