(self, x)
| 85 | return mask |
| 86 | |
| 87 | def __call__(self, x): |
| 88 | # Extract some shapes |
| 89 | B, N = x.shape |
| 90 | eos_tokens = x.argmax(-1) |
| 91 | |
| 92 | # Compute the embeddings |
| 93 | x = self.token_embedding(x) |
| 94 | x = x + self.position_embedding.weight[:N] |
| 95 | |
| 96 | # Compute the features from the transformer |
| 97 | mask = self._get_mask(N, x.dtype) |
| 98 | hidden_states = [] |
| 99 | for l in self.layers: |
| 100 | x = l(x, mask) |
| 101 | hidden_states.append(x) |
| 102 | |
| 103 | # Apply the final layernorm and return |
| 104 | x = self.final_layer_norm(x) |
| 105 | last_hidden_state = x |
| 106 | |
| 107 | # Select the EOS token |
| 108 | pooled_output = x[mx.arange(len(x)), eos_tokens] |
| 109 | if "text_projection" in self: |
| 110 | pooled_output = self.text_projection(pooled_output) |
| 111 | |
| 112 | return CLIPOutput( |
| 113 | pooled_output=pooled_output, |
| 114 | last_hidden_state=last_hidden_state, |
| 115 | hidden_states=hidden_states, |
| 116 | ) |
nothing calls this directly
no test coverage detected