PyTorch Streaming Softmax

Task

Implement streaming_softmax(logits: torch.Tensor, chunk_size: int) -> torch.Tensor. Debug a chunked PyTorch softmax: replace the broken per-chunk normalization with a numerically stable streaming scan that keeps one global online maximum and denominator.

Requirements

• The starter implementation is intentionally buggy and should be repaired without changing the function signature.
• logits is a non-empty 1D tensor.
• chunk_size is a positive integer; the last chunk may be smaller.
• Do not exponentiate unshifted logits.
• Maintain a running maximum and rescale the previous denominator whenever a later chunk has a larger maximum.
• Return a tensor with the same shape as logits whose values sum to 1.
• Preserve logits dtype and device in the returned tensor.

Example

streaming_softmax(torch.tensor([1.0, 2.0, 3.0, 4.0]), 2)
# tensor([0.0321, 0.0871, 0.2369, 0.6439])