PyTorch Coding Practice Problems

94 Leetcode-style machine learning problems solved in PyTorch and NumPy, from single operators up to full transformer and inference stacks.

Easy

• Transformer Positional Encoding Implement absolute positional encoding for transformer models, adding position information to token embeddings so the model can understand sequence order.
• Batch Normalization Implement batch normalization layer that normalizes activations across mini-batches to stabilize and accelerate neural network training.
• Build Binary Classifier with Cross Entropy Loss Build a binary classifier from scratch using binary cross-entropy loss and gradient descent optimization.
• Binary Quantization and Bit Packing Encode vector signs as packed bytes with explicit bit order and zero padding.
• Score Allowed Candidate Tokens Project final hidden states onto allowed vocabulary rows and return masked candidate logits.
• Cross Entropy Loss Implement cross entropy loss functions for classification tasks, from Binary Cross Entropy (BCE) for binary classification to multi-class Cross Entropy for neural networks.
• Implementing Leaky ReLU with Manual Autograd Build a custom Leaky ReLU activation function with manual autograd (backward pass) to understand gradient computation in neural networks.
• Data Query API Design a data query API system that efficiently handles filtering, sorting, and pagination of large datasets.
• DSpark: Confidence and Prefix Survival Turn per-position confidence scores into prefix survival probabilities and the expected number of accepted draft tokens.
• Gaussian Blur Kernel Implement a function that generates a 2D Gaussian blur kernel of size k x k, used in image processing to smooth images by reducing noise and detail.
• Gated DeltaNet: Projections and Head Shapes Project hidden states into independent query, key, value, retention, write, and output-gate branches.
• Gated DeltaNet: Q/K L2 Normalization Normalize each query and key head over its feature axis with a stable L2 denominator.
• Gated DeltaNet: Retention and Write Gates Convert projected logits into per-head retention and write strengths.
• Mask Logits with a Grammar State Mask tokens that have no valid transition from each request’s grammar state.
• Decode with Greedy Sampling Implement greedy sampling for language model text generation, selecting the highest probability token at each step.
• Route Requests by Replica Load Route arrivals sequentially to available replicas using capacity-normalized outstanding work.
• KL Divergence Penalty for RLHF Calculate KL divergence penalty for RLHF (Reinforcement Learning from Human Feedback) to constrain policy updates during training.
• Layer Norm Implement layer normalization that normalizes activations across features for each sample, commonly used in transformer architectures.
• Looped Transformer: Shared-Weight Unrolling Reuse one block across depth and return the hidden-state history.
• Matryoshka Embedding Search Search trained embedding prefixes with deterministic dot-product ranking.
• Max Pooling Implement 2D max pooling operation used in convolutional neural networks to downsample feature maps while preserving important spatial features.
• Accumulate and Average Replica Gradients Accumulate microbatch gradients on each data-parallel replica and average the matching parameter shard by sample count.
• MoE: Top-k Expert Routing Select expert IDs per token and normalize only their selected routing scores.
• Pairwise Cosine Similarity Compute a pairwise cosine similarity matrix in PyTorch with safe L2 normalization.
• Posting List Intersection Intersect sorted posting lists to find IDs shared by every list.
• Quantization: INT8 Round Trip Round, saturate, and reconstruct values with a supplied symmetric quantization scale.
• RMSNorm: Root Mean Square Layer Norm Implement RMSNorm (Root Mean Square Layer Normalization), a simpler and faster alternative to LayerNorm used in models like LLaMA.
• Calculate Self-Attention Calculate self-attention mechanism that allows each token to attend to all other tokens in the sequence, the core of transformer models.
• Numerically Stable Softmax Implement the numerically stable softmax function that converts a list of logits into a probability distribution, using the max-subtraction trick to avoid overflow.
• Split a Dataset by Source Family Build reproducible dataset partitions without separating source families or contrast pairs.
• Temporal Difference (TD) Error for Action-Value Function Implement single-step Temporal Difference (TD) learning to update the Q-function, the most widely used method for estimating action-value functions in reinforcement learning.
• Temperature Sampling with Exponential Noise Select tokens using temperature-scaled scores and supplied exponential-noise draws.
• Triplet Margin Loss Implement a mean triplet hinge loss in PyTorch using squared Euclidean distances.

Medium

• Optimizers: From Momentum to Adam Learn optimizer evolution from SGD with Momentum to RMSProp to Adam. Implement each optimizer step-by-step to understand adaptive learning rates and momentum in deep learning.
• AdamW Optimizer Implement AdamW optimizer with weight decay decoupling, the most popular optimizer for training large language models and transformers.
• Implement Bagging Classifier with Decision Stumps Implement bagging ensemble method with decision stumps as weak learners to improve classification accuracy through voting.
• Basic KV Cache: Prefill and Decode Implement KV-cache for efficient transformer inference, storing key-value pairs to avoid recomputing attention for previous tokens.
• Binary Vector Scoring with Lookup Tables Build signed-sum lookup tables and score binary vectors against a floating-point query.
• Bradley Terry Ranking Implement Bradley-Terry model for pairwise comparison ranking, commonly used in RLHF for preference learning.
• Build KD Tree Build a KD-tree data structure for efficient nearest neighbor search in multi-dimensional space.
• Validate Counterfactual Evidence Pairs Keep training pairs that change one known fact, flip the answer, and require the supporting evidence.
• DAPO: Clip-Higher Policy Optimization Loss Implement DAPO (Dynamic and Adaptive Policy Optimization) loss function, a reinforcement learning technique that uses a Clip-Higher mechanism to prevent entropy collapse during chain-of-thought reasoning in language models.
• Direct Preference Optimization Loss Implement DPO (Direct Preference Optimization) loss, a simpler alternative to RLHF that directly optimizes language models from preference data without requiring a separate reward model.
• DSpark: Markov Draft Head Draft several tokens from one set of parallel backbone logits by adding a low-rank bias conditioned on the previously drafted token.
• Exact Vector Reranking Rescore a supplied candidate set using full vectors and return deterministic top results.
• Fit Temperature Scaling by Log Loss Choose a positive calibration temperature by stable held-out log loss over allowed candidates.
• Gated DeltaNet: Recurrent Sequence Carry gated delta memory through a causal sequence and resume from a saved state.
• Gated DeltaNet: One Memory Update Decay memory, correct its prediction, write a rank-one update, and read the updated state.
• Gated DeltaNet: Causal Depthwise Convolution Mix local token history independently in each channel, then apply SiLU.
• Gated DeltaNet: Gated Output Projection Apply learned per-head RMS normalization, a SiLU output gate, and the final projection.
• Group Relative Policy Optimization Loss Implement GRPO (Group Relative Policy Optimization) loss for reinforcement learning, using group-relative rewards to improve training stability.
• Pack Prefill and Decode Batches Pack scheduled token slices and map their logical positions to physical KV-cache slots.
• Measure Token Latency and SLO Goodput Aggregate first-token latency, token gaps, throughput, and request goodput from serving traces.
• Inverted File (IVF) Search Build centroid-assigned inverted lists and probe selected clusters for candidate IDs.
• K-Means Clustering Implement K-means clustering algorithm that partitions data into K clusters by minimizing within-cluster variance.
• Knowledge Distillation Loss Implement knowledge distillation loss that trains a smaller student model to mimic a larger teacher model using soft probability distributions.
• Linear Regression: From OLS to Lasso Master linear regression from closed-form OLS solution through gradient descent to regularization techniques (Ridge L2 and Lasso L1) for preventing overfitting.
• Longest Uninterrupted Work Time Find the longest uninterrupted time slot in a schedule by analyzing meeting intervals and gaps.
• Looped Transformer: Adaptive Early Exit Stop each sequence independently when its hidden-state update is small enough.
• From LoRA to QLoRA Learn parameter-efficient fine-tuning from vanilla LoRA to quantized QLoRA. Implement low-rank adaptation and 4-bit quantization for efficient LLM fine-tuning.
• Max Pooling with Coordinates Implement 2D max pooling that returns the coordinates of the maximum elements for unpooling operations.
• Schedule Pipeline Microbatches Build the one-forward-one-backward pipeline schedule and report each stage's warmup depth and peak stored activations.
• Schedule Decode and Prefill Within a Token Budget Share an engine-step token budget between decode requests and bounded prefill chunks.
• MoE: Dispatch and Weighted Combine Batch selected tokens by expert and scatter weighted expert outputs back to token order.
• MoE: Router Load-Balancing Loss Compute a masked auxiliary router loss from hard assignments and full expert probabilities.
• Beam Search with nanoGPT Implement beam search decoding for language model text generation with KV-cache reordering. Maintain multiple candidate sequences in parallel to find higher-quality outputs than greedy decoding.
• NT-Xent Contrastive Loss Implement the SimCLR NT-Xent objective in PyTorch with positive-pair targets, self-pair masking, and stable temperature-scaled logits.
• PagedAttention: KV Cache Block Manager Implement PagedAttention block manager for efficient GPU memory management by organizing KV-cache into fixed-size blocks.
• PagedAttention: Paged Attention Score Calculate attention scores using PagedAttention, which stores key-value pairs in non-contiguous memory blocks for efficient serving.
• PagedAttention: Shared Prefix Mapping Implement shared prefix mapping for PagedAttention to deduplicate common prompt prefixes and save memory.
• Pointer Readout over Options Compute projected decision-to-option scores with biases, pointer-width scaling, and candidate masking.
• Coordinate Prefill-to-Decode KV Handoff Simulate a KV-transfer lifecycle that handles acknowledgments, cancellation, and late events.
• QAT: Quantization Aware Training Implement Quantization-Aware Training (QAT) to prepare neural networks for efficient low-precision deployment, from fake quantization with STE to complete QAT network training.
• Quantization: Per-Tensor and Per-Channel Scales Choose robust AbsMax scales and compare tensor-wide versus output-channel quantization.
• Quantization: Weight-Only Linear Layer Apply grouped quantized weights while preserving floating-point activations and batched shapes.
• RoPE: Precompute Rotary Embeddings Precompute RoPE (Rotary Position Embeddings) frequencies and rotation matrices for efficient positional encoding in transformers.
• Sliding Window Attention Implement sliding window attention mechanism that limits attention span to a local window, reducing computational complexity for long sequences.
• Stateful Batch Normalization Implement stateful batch normalization that tracks running statistics across batches for inference mode.
• Column and Row Tensor Parallel Linear Simulate tensor-parallel linear projections and expose each shard contribution on a single GPU.
• Decode with TopK Sampling Implement top-k sampling for language model text generation, sampling from the k most probable tokens to enable more diverse outputs than greedy decoding.
• Decode with TopP Sampling Implement nucleus (top-p) sampling for language model decoding, restricting each draw to the smallest set of tokens whose probabilities reach p.
• Unbiased GRPO Optimization: Dr. GRPO Implement Dr. GRPO (Unbiased Group Relative Policy Optimization) that removes baseline bias for more stable reinforcement learning.

Hard

• Implement an AI ToDo List System Design and implement an AI-powered todo list system with intelligent task prioritization, scheduling, and natural language processing.
• Applying RoPE to Q and K in Causal Attention Apply RoPE (Rotary Position Embeddings) to query and key matrices in causal attention for position-aware attention computation.
• DSpark: Hardware-Aware Prefix Scheduling Trim a speculative draft to the prefix length that maximizes output tokens per second, combining prefix survival with the engine's measured batch-speed table.
• Build Causal Grouped-Query Attention (GQA) Build causal grouped-query attention (GQA) that shares key-value heads across multiple query heads to reduce memory usage while maintaining quality.
• Continuous Batching and Preemption Plan one prefill-first engine step while allocating blocks and preempting decode requests.
• Tensor-Parallel Forward and Backward Collectives Simulate the conjugate copy/all-reduce operators of a tensor-parallel MLP and rebuild the dense gradients from shard gradients.
• Vocabulary-Parallel Embedding and Cross-Entropy Mask out-of-range embedding lookups per rank and reduce sharded logits into a numerically stable cross-entropy loss.
• Build Causal Multi-Head Attention for nanoGPT Build causal multi-head attention for nanoGPT, splitting attention into multiple heads to capture different types of relationships.
• Build Causal Multi-Query Attention (MQA) Build causal multi-query attention (MQA) that uses a single key-value head shared across all query heads for faster inference.
• ReLU²: Thresholded Activation for MoE Implement ReLU² (squared ReLU) activation function used as a thresholding mechanism in Mixture-of-Experts (MoE) gating.
• Speculative Decoding: Rejection Sampling Implement speculative decoding with rejection sampling to speed up language model inference by using a smaller draft model.
• Accept or Reject Draft Tokens Verify draft tokens and sample a corrected continuation using deterministic uniform draws.