Explicit Expansion in MHAEinsum is Redundant

[jestinepaul]

# MHAEinsum - explicit expansion (redundant)
mask = self.mask[:n, :n].unsqueeze(0).unsqueeze(1).expand(b, self.num_heads, n, n)
scores = scores.masked_fill(mask.bool(), -torch.inf)

# Ch03_MHA - relies on broadcasting (efficient)
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
attn_scores.masked_fill_(mask_bool, -torch.inf)

PyTorch's masked_fill_ automatically broadcasts the mask to match the target tensor's shape. Since scores has shape (b, num_heads, n, n) and the mask has shape (n, n), PyTorch will automatically broadcast the 2D mask across the batch and head dimensions.

The Ch03_MHA approach is more efficient because:

1. Less memory usage - no need to create a 4D tensor
2. Faster execution - no explicit expansion operations
3. Cleaner code - leverages PyTorch's built-in broadcasting

The MHAEinsum expansion step can be simplified to match Ch03_MHA's approach without any functional difference.

[rasbt]

Thanks for the suggestion! I agree, this is more efficient, and definitely more elegant. I don't know why I had issues with it before. I'm updating it via #775