def __call__(
      self,
      encoded_input: Array,
      mlm_target_positions: Array,
      shared_embedding: Array,
  ) -> Array:
    """Perform masked language modeling scoring.

    Args:
      encoded_input: [bsz, n_tokens, hidden_size].
      mlm_target_positions: [bsz, max_mlm_targets] positions of mlm targets in
        passage.
      shared_embedding: [vocab_size, hidden_size] word embedding array, shared
        with initial embedding.

    Returns:
      Array of masked language modeling logits.
    """

    target_encodings = jut.matmul_slice(encoded_input, mlm_target_positions)
    target_encodings = self.dense(target_encodings)
    target_encodings = nn.gelu(target_encodings)
    target_encodings = self.layer_norm(target_encodings)

    mlm_logits = self.embedding_dense.apply(
        {'params': {
            'kernel': shared_embedding.T
        }}, target_encodings)
    mlm_logits = mlm_logits + self.bias

    return mlm_logits
Perform masked language modeling scoring.

    Args:
      encoded_input: [bsz, n_tokens, hidden_size].
      mlm_target_positions: [bsz, max_mlm_targets] positions of mlm targets in
        passage.
      shared_embedding: [vocab_size, hidden_size] word embedding array, shared
        with initial embedding.

    Returns:
      Array of masked language modeling logits.
    
language/mentionmemory/modules/mlm_layer.py