List of all items

Structs

• activation::ELU
• activation::GELU
• activation::LeakyReLU
• activation::Mish
• activation::ReLU
• activation::SELU
• activation::SiLU
• activation::Sigmoid
• activation::Softmax
• activation::Softplus
• activation::Swish
• activation::Tanh
• adversarial::AdversarialAttack
• adversarial::AdversarialTrainer
• adversarial::AdversarialTrainingConfig
• adversarial::AttackConfig
• adversarial::RandomizedSmoothing
• attention::CrossAttention
• attention::MultiHeadAttention
• bert::BertConfig
• bert::BertEmbeddings
• bert::BertForMaskedLM
• bert::BertForSequenceClassification
• bert::BertForTokenClassification
• bert::BertModel
• bert::BertOutput
• bert::BertPooler
• clip::CLIP
• clip::CLIPConfig
• clip::CLIPTextConfig
• clip::CLIPTextEncoder
• clip::CLIPTextLayer
• clip::CLIPVisionConfig
• clip::CLIPVisionEncoder
• clip::LayerNorm
• clip::MLP
• clip::MultiHeadAttention
• conv::Conv1d
• conv::Conv2d
• conv::Conv3d
• conv::TransposeConv2d
• curriculum_learning::CurriculumConfig
• curriculum_learning::CurriculumLearning
• curriculum_learning::CurriculumStats
• curriculum_learning::DifficultyScorer
• curriculum_learning::ScoredSample
• differential_privacy::DPConfig
• differential_privacy::DPSGDOptimizer
• differential_privacy::LocalDP
• differential_privacy::PATEEnsemble
• differential_privacy::PrivacyAccountant
• diffusion::DDPM
• diffusion::DiffusionConfig
• diffusion::NoiseScheduler
• diffusion::ResidualBlock
• diffusion::UNet
• distributed::DataParallel
• distributed::DistributedConfig
• distributed::DistributedDataParallel
• distributed::GradientAccumulator
• distributed::ModelParallel
• distributed::PipelineParallel
• dropout::Dropout
• dropout::Dropout2d
• embedding::Embedding
• embedding::TokenPositionEmbedding
• federated::DifferentialPrivacy
• federated::FederatedClient
• federated::FederatedServer
• federated::SecureAggregation
• flash_attention::FlashAttention
• flash_attention::FlashAttentionConfig
• gnn::GATLayer
• gnn::GCNLayer
• gnn::Graph
• gnn::GraphSAGELayer
• gnn::MPNNLayer
• gpt::GPTConfig
• gpt::GPTEmbeddings
• gpt::GPTForCausalLM
• gpt::GPTForSequenceClassification
• gpt::GPTModel
• gradient_checkpointing::CheckpointConfig
• gradient_checkpointing::CheckpointManager
• gradient_checkpointing::CheckpointStats
• gradient_checkpointing::CheckpointedLayer
• gradient_checkpointing::CheckpointedSequential
• inference::BatchInference
• inference::FusedOp
• inference::InferenceConfig
• inference::InferenceOptimizer
• inference::InferenceSession
• knowledge_distillation::DistillationConfig
• knowledge_distillation::DistillationStats
• knowledge_distillation::KnowledgeDistillation
• linear::Linear
• llama::LLaMAAttention
• llama::LLaMAConfig
• llama::LLaMADecoderLayer
• llama::LLaMAForCausalLM
• llama::LLaMAModel
• llama::RMSNorm
• llama::RotaryEmbedding
• llama::SwiGLU
• lora::LoRAConfig
• lora::LoRALinear
• lora::QLoRAConfig
• lora::QLoRALinear
• mesh::Mesh
• mesh::MeshConv
• mesh::MeshEncoder
• mesh::MeshPool
• mesh::MeshUtils
• mixed_precision::AutocastContext
• mixed_precision::GradScaler
• mixed_precision::MixedPrecisionConfig
• mixture_of_experts::Expert
• mixture_of_experts::MixtureOfExperts
• mixture_of_experts::MoEConfig
• mixture_of_experts::MoEStats
• mixture_of_experts::Router
• module::Sequential
• nerf::NeRF
• nerf::NeRFConfig
• nerf::NeRFMLP
• nerf::PositionalEncoder
• nerf::RaySampler
• nerf::VolumeRenderer
• norm::BatchNorm1d
• norm::BatchNorm2d
• norm::GroupNorm
• norm::InstanceNorm
• norm::LayerNorm
• onnx::ONNXGraph
• onnx::ONNXModel
• onnx::ONNXNode
• onnx::ONNXTensor
• point_cloud::FarthestPointSampler
• point_cloud::KNNGrouper
• point_cloud::PointNet
• point_cloud::PointNetBackbone
• point_cloud::PointNetConfig
• point_cloud::STN3d
• pooling::AdaptiveAvgPool2d
• pooling::AvgPool2d
• pooling::GlobalAvgPool2d
• pooling::GlobalMaxPool2d
• pooling::MaxPool2d
• prompt_tuning::LayerPrefix
• prompt_tuning::PTuningV2
• prompt_tuning::PrefixTuning
• prompt_tuning::PrefixTuningConfig
• prompt_tuning::PromptTuning
• prompt_tuning::PromptTuningConfig
• quantization::DynamicQuantization
• quantization::QuantizationAwareTraining
• quantization::QuantizationConfig
• quantization::QuantizedTensor
• ring_attention::RingAttention
• ring_attention::RingAttentionConfig
• ring_attention::StripedRingAttention
• rl::ActorCriticAgent
• rl::DQNAgent
• rl::Experience
• rl::PPOAgent
• rl::PolicyNetwork
• rl::QNetwork
• rl::REINFORCEAgent
• rl::ReplayBuffer
• rl::ValueNetwork
• rnn::GRU
• rnn::GRUCell
• rnn::LSTM
• rnn::LSTMCell
• serialization::ModelCheckpoint
• serialization::ModelMetadata
• t5::T5Config
• t5::T5Decoder
• t5::T5Embeddings
• t5::T5Encoder
• t5::T5ForConditionalGeneration
• t5::T5ForSequenceClassification
• t5::T5Model
• t5::T5Output
• transformer::FeedForward
• transformer::PositionalEncoding
• transformer::RotaryEmbedding
• transformer::TransformerDecoderLayer
• transformer::TransformerEncoder
• transformer::TransformerEncoderLayer
• vision_transformer::PatchEmbedding
• vision_transformer::ViTConfig
• vision_transformer::VisionTransformer
• zero_optimizer::ParameterPartition
• zero_optimizer::ZeRoConfig
• zero_optimizer::ZeRoOptimizer
• zero_optimizer::ZeRoStats

Enums

• adversarial::AttackType
• curriculum_learning::CurriculumStrategy
• curriculum_learning::DifficultyMetric
• curriculum_learning::PacingFunction
• diffusion::BetaSchedule
• distributed::DistributedBackend
• federated::AggregationStrategy
• gnn::AggregatorType
• gradient_checkpointing::CheckpointStrategy
• knowledge_distillation::DistillationMethod
• lora::QuantType
• mixed_precision::PrecisionMode
• mixture_of_experts::RoutingStrategy
• onnx::ONNXAttribute
• onnx::ONNXDataType
• prompt_tuning::PromptInitStrategy
• quantization::QuantizationScheme
• transformer::Activation
• zero_optimizer::ZeRoStage

Traits

• module::Module

Functions

• attention::scaled_dot_product_attention
• gradient_checkpointing::estimate_memory_savings
• inference::warmup_model
• init::constant
• init::kaiming_normal
• init::kaiming_uniform
• init::normal
• init::ones
• init::uniform
• init::xavier_normal
• init::xavier_uniform
• init::zeros
• loss::binary_cross_entropy
• loss::binary_cross_entropy_with_logits
• loss::contrastive_loss
• loss::cosine_embedding_loss
• loss::cross_entropy
• loss::focal_loss
• loss::huber_loss
• loss::l1_loss
• loss::mse_loss
• loss::nll_loss
• loss::smooth_l1_loss
• loss::triplet_margin_loss
• mixed_precision::to_full_precision
• mixed_precision::to_half_precision
• onnx::onnx_to_tensor
• onnx::tensor_to_onnx
• serialization::load_model
• serialization::save_model