Re-exports§
pub use activation::Activation;pub use conv1d::Conv1d;pub use conv2d::Conv2d;pub use dropout::Dropout;pub use embedding::Embedding;pub use groupnorm::GroupNorm;pub use layernorm::LayerNorm;pub use linear::Linear;pub use linear::MaybeQuantLinear;pub use linear::QuantLinear;pub use lora::LoraLinear;pub use loss::contrastive_loss;pub use loss::cross_entropy_loss;pub use loss::cross_entropy_loss_smooth;pub use loss::focal_loss;pub use loss::kl_div_loss;pub use loss::mse_loss;pub use mla::Mla;pub use mla::MlaConfig;pub use module::Module;pub use module::StateDict;pub use module::TrainMode;pub use moe::Expert;pub use moe::MoeLayer;pub use moe::MoeLayerConfig;pub use moe::MoeOutput;pub use moe::MoeRouter;pub use moe::MoeRouterConfig;pub use moe::RouterOutput;pub use rmsnorm::RmsNorm;pub use rope::RoPE;pub use stochastic_depth::StochasticDepth;pub use var_builder::VarBuilder;pub use varmap::Init;pub use varmap::VarMap;pub use weight::Weight;
Modules§
- activation
- Activation function enum for configurable model architectures
- conv1d
- 1D convolution layer with autograd support
- conv2d
- 2D convolution layer with autograd support
- dropout
- Dropout regularization layer
- embedding
- Embedding layer — lookup table for token embeddings
- groupnorm
- Group Normalization module
- layernorm
- Layer Normalization module
- linear
- Linear and quantized linear layers
- lora
- LoRA (Low-Rank Adaptation) layer.
- loss
- mla
- Multi-Head Latent Attention (MLA) module
- module
- Neural network module traits for parameter access and serialization.
- moe
- rmsnorm
- RMS Normalization module
- rope
- RoPE (Rotary Position Embedding) module
- stochastic_
depth - Stochastic depth (drop path) regularization.
- var_
builder - VarBuilder: scoped access to weights in a VarMap.
- varmap
- weight
- Weight enum for storing standard or quantized tensors in VarMap