(self, x: mx.array)
| 202 | self.final_layer_norm = nn.LayerNorm(config.hidden_size) |
| 203 | |
| 204 | def __call__(self, x: mx.array) -> CLIPTextOutput: |
| 205 | B, N = x.shape |
| 206 | eot_tokens = mx.argmax(x, axis=-1) |
| 207 | x = self.embeddings(x) |
| 208 | mask = nn.MultiHeadAttention.create_additive_causal_mask(N, x.dtype) |
| 209 | for l in self.encoder.layers: |
| 210 | x = l(x, mask) |
| 211 | last_hidden_state = self.final_layer_norm(x) |
| 212 | pooler_output = last_hidden_state[mx.arange(B), eot_tokens] |
| 213 | |
| 214 | return CLIPTextOutput( |
| 215 | pooler_output=pooler_output, last_hidden_state=last_hidden_state |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | class VisionEmbeddings(nn.Module): |
nothing calls this directly
no test coverage detected