import torch
import torch.nn as nn
import torch.nn.functional as F

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.relu(self.w_1(x))
        h = self.dropout(h)
        return self.w_2(h)

if __name__ == '__main__':
    D_MODEL = 512
    D_FF = 2048
    BATCH_SIZE = 4
    SEQ_LEN = 60
    DROPOUT = 0.1

    ffn = PositionwiseFeedForward(D_MODEL, D_FF, DROPOUT)
    dummy_input = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
    output = ffn(dummy_input)
    
    print(f"入力形状: {dummy_input.shape}")
    print(f"出力形状: {output.shape}")
    assert output.shape == (BATCH_SIZE, SEQ_LEN, D_MODEL)

# Attention Is All You Need
# 3.3 Position-wise Feed-Forward Networks
# PositionwiseFeedForward.py
# 第2部 第7章