(
self, input_ids: mx.array, token_type_ids: mx.array = None
)
| 71 | self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 72 | |
| 73 | def __call__( |
| 74 | self, input_ids: mx.array, token_type_ids: mx.array = None |
| 75 | ) -> mx.array: |
| 76 | words = self.word_embeddings(input_ids) |
| 77 | position = self.position_embeddings( |
| 78 | mx.broadcast_to(mx.arange(input_ids.shape[1]), input_ids.shape) |
| 79 | ) |
| 80 | |
| 81 | if token_type_ids is None: |
| 82 | # If token_type_ids is not provided, default to zeros |
| 83 | token_type_ids = mx.zeros_like(input_ids) |
| 84 | |
| 85 | token_types = self.token_type_embeddings(token_type_ids) |
| 86 | |
| 87 | embeddings = position + words + token_types |
| 88 | return self.norm(embeddings) |
| 89 | |
| 90 | |
| 91 | class Bert(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected