| 13 | |
| 14 | @dataclass |
| 15 | class CLIPOutput: |
| 16 | # The last_hidden_state indexed at the EOS token and possibly projected if |
| 17 | # the model has a projection layer |
| 18 | pooled_output: Optional[mx.array] = None |
| 19 | |
| 20 | # The full sequence output of the transformer after the final layernorm |
| 21 | last_hidden_state: Optional[mx.array] = None |
| 22 | |
| 23 | # A list of hidden states corresponding to the outputs of the transformer layers |
| 24 | hidden_states: Optional[List[mx.array]] = None |
| 25 | |
| 26 | |
| 27 | class CLIPEncoderLayer(nn.Module): |