Implement focal loss for class-imbalanced binary classification from basic tensor operations, keeping it numerically stable in logit space.

Focal loss scales cross-entropy by (1 − p_t) to the power gamma, so well-classified examples add little and training focuses on hard ones; alpha rebalances the classes. Work from logits throughout.

import torch

def focal_loss(z, y, alpha=0.25, gamma=2.0):
    # z: logits, y: float 0/1 targets, same shape
    # stable BCE: max(z,0) - z*y + log(1+exp(-|z|))
    ce = (z.clamp(min=0) - z * y
          + torch.log1p(torch.exp(-z.abs())))
    p_t = torch.exp(-ce)
    a_t = alpha * y + (1 - alpha) * (1 - y)
    return (a_t * (1 - p_t) ** gamma * ce).mean()

Interviewers check you never take the log of a sigmoid output, that p_t comes from exp of minus the loss, and that gamma 0 recovers alpha-weighted cross-entropy.