feat: support cross_entropy

Author: RayZhao1998

cs336_basics/cross_entropy.py

@@ -0,0 +1,19 @@
+from jaxtyping import Float, Int
+from torch import Tensor, exp, log
+
+def cross_entropy(inputs: Float[Tensor, "batch_size vocab_size"], targets: Int[Tensor, "batch_size"]) -> Float[Tensor, ""]:
+    max_logits = inputs.max(dim=1, keepdim=True).values # [batch_size, 1]
+    inputs_stable = inputs - max_logits # [batch_size, vocab_size]
+
+    target_logits = inputs_stable.gather(dim=1, index=targets.unsqueeze(dim=1))
+
+    exp_inputs = exp(inputs_stable)
+    sum_exp = exp_inputs.sum(dim=1, keepdim=True)
+    log_sum_exp = log(sum_exp)
+
+    sample_losses = -target_logits + log_sum_exp
+
+    avg_loss = sample_losses.flatten().mean()
+
+    return avg_loss
+

tests/adapters.py

@@ -21,6 +21,7 @@
 from cs336_basics.multihead_self_attention import MultiHeadSelfAttention, MultiHeadSelfAttentionWithRoPE
 from cs336_basics.transformer_block import TransformerBlock
 from cs336_basics.transformer import Transformer
+from cs336_basics.cross_entropy import cross_entropy
 
 def run_linear(
     d_in: int,
@@ -511,7 +512,7 @@ def run_cross_entropy(inputs: Float[Tensor, " batch_size vocab_size"], targets:
     Returns:
         Float[Tensor, ""]: The average cross-entropy loss across examples.
     """
-    raise NotImplementedError
+    return cross_entropy(inputs, targets)
 
 
 def run_gradient_clipping(parameters: Iterable[torch.nn.Parameter], max_l2_norm: float) -> None: