Implement a Transformer Block

Problem

Implement a decoder-style transformer block in PyTorch from scratch. Write a module that takes an input tensor x of shape (batch_size, seq_len, d_model) and returns an output tensor of the same shape after applying causal multi-head self-attention, a position-wise feed-forward network, residual connections, and layer normalization. Do not use torch.nn.MultiheadAttention; implement the attention logic directly with linear projections and tensor reshaping.

Constraints

• 1 <= batch_size <= 32
• 1 <= seq_len <= 512
• 1 <= d_model <= 1024
• d_model % num_heads == 0
• Use PyTorch tensor operations for attention, masking, and projection