| 102 | |
| 103 | |
| 104 | class EncoderLayer(nn.Module): |
| 105 | def __init__(self, config: VisionConfig): |
| 106 | super().__init__() |
| 107 | self.embed_dim = config.hidden_size |
| 108 | self.self_attn = Attention( |
| 109 | config.hidden_size, config.num_attention_heads, bias=True |
| 110 | ) |
| 111 | self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 112 | self.mlp = MLP(config) |
| 113 | self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 114 | |
| 115 | def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array: |
| 116 | y = self.layer_norm1(x) |
| 117 | y = self.self_attn(y, y, y, mask) |
| 118 | x = x + y |
| 119 | y = self.layer_norm2(x) |
| 120 | y = self.mlp(y) |
| 121 | return x + y |
| 122 | |
| 123 | |
| 124 | class Encoder(nn.Module): |