feat: support softmax function

Author: RayZhao1998

cs336_basics/softmax.py

@@ -0,0 +1,10 @@
+import torch
+from jaxtyping import Float
+
+def softmax(tensor: Float[torch.Tensor, "..."], dim: int) -> Float[torch.Tensor, "..."]:
+    max = tensor.max(dim=dim, keepdim=True).values
+    tensor_shifted = tensor - max
+    exp_tensor = torch.exp(tensor_shifted)
+    exp_sum = exp_tensor.sum(dim=dim, keepdim=True)
+    result = exp_tensor / exp_sum
+    return result

tests/adapters.py

@@ -16,6 +16,7 @@
 from cs336_basics.rmsnorm import RMSNorm
 from cs336_basics.swiglu import SwiGLU
 from cs336_basics.rope import RoPE
+from cs336_basics.softmax import softmax
 
 def run_linear(
     d_in: int,
@@ -446,7 +447,7 @@ def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, "
         Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
         softmax normalizing the specified `dim`.
     """
-    raise NotImplementedError
+    return softmax(in_features, dim)
 
 
 def run_cross_entropy(inputs: Float[Tensor, " batch_size vocab_size"], targets: Int[Tensor, " batch_size"]) -> Float[Tensor, ""]: