Encoder出力（文脈表現）Contextual Representation

from dataclasses import dataclass
from typing import Optional
import torch
from torch import Tensor, nn
 
 
@dataclass
class EncoderOutput:
    encoder_states: Tensor            # [B, S, D] 文脈込みの表現
    encoder_mask: Optional[Tensor]    # [B, S]    パディング位置
 
 
class TransformerEncoder(nn.Module):
    """Encoder 層を N 個積み重ねたスタック。"""
 
    def __init__(self, layers: nn.ModuleList) -> None:
        super().__init__()
        # 各層は別々のパラメータを持つ。同じ層を N 回ループするのではない。
        self.layers = layers
 
    def forward(
        self,
        hidden_states: Tensor,
        encoder_mask: Optional[Tensor] = None,
    ) -> EncoderOutput:
        """
        Args:
            hidden_states: [B, S, D]。位置情報を加算済みの埋め込み。
            encoder_mask: [B, S]。1 が実トークン、0 がパディング。
 
        Returns:
            encoder_states [B, S, D]。Decoder が Key/Value として参照する。
        """
        for layer in self.layers:
            hidden_states = layer(hidden_states, encoder_mask)  # 形は [B, S, D] のまま
        return EncoderOutput(encoder_states=hidden_states, encoder_mask=encoder_mask)
 
 
# 最小使用例（層の中身はダミー）
B, S, D = 1, 8, 512
identity_layers = nn.ModuleList([nn.Identity() for _ in range(6)])
 
 
class _DummyLayer(nn.Module):
    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        return x
 
 
encoder = TransformerEncoder(nn.ModuleList([_DummyLayer() for _ in range(6)]))
out = encoder(torch.randn(B, S, D), torch.ones(B, S, dtype=torch.long))
 
print(out.encoder_states.shape)   # torch.Size([1, 8, 512]) == [B, S, D]
print(out.encoder_mask.shape)     # torch.Size([1, 8])      == [B, S]