Softmax

Numerical stability, exp/log tricks

Easy Fundamentals

Problem Description

Implement the Softmax function from scratch.

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$$

Signature

def my_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor: ...

Rules

• Do NOT use torch.softmax, F.softmax, or torch.nn.Softmax

• Must be numerically stable (hint: subtract max before exp)

Example

Input: tensor([1., 2., 3.]) Output: tensor([0.0900, 0.2447, 0.6652]) # sums to 1.0

Template

Implement the function below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

def my_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    pass  # Replace this

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
x = torch.tensor([1.0, 2.0, 3.0])
print("Output:", my_softmax(x, dim=-1))
print("Sum:   ", my_softmax(x, dim=-1).sum())  # should be ~1.0
print("Ref:   ", torch.softmax(x, dim=-1))

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

def my_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    x_max = x.max(dim=dim, keepdim=True).values
    e_x = torch.exp(x - x_max)
    return e_x / e_x.sum(dim=dim, keepdim=True)

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("softmax")

Key Concepts

Numerical stability, exp/log tricks

ReLU Cross-Entropy Loss

Softmax

Description Template Test Solution Tips