class_docstr: Head used for masked language modelling in BERT.
  
  Attributes:
    dtype: Data type.
class_docstr_tok: ["Head","used","for","masked","language","modelling","in","BERT",".","Attributes",":","dtype",":","Data","type","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      masked_lm_positions: jnp.ndarray,
      word_embeddings: jnp.ndarray,
      *,
      train: bool,
  ) -> jnp.ndarray:
      batch_size, length, hidden_size = x.shape
      x = x.reshape((-1, hidden_size))
  
      offsets = np.arange(batch_size)[:, None] * length
      masked_lm_positions = masked_lm_positions + offsets
      masked_lm_positions = masked_lm_positions.ravel()
  
      x = jnp.take(x, masked_lm_positions, axis=0)
      x = x.reshape((batch_size, -1, hidden_size))
  
      vocab_size, embedding_width = word_embeddings.shape
      kernel_init = bert_truncated_normal_initializer()
  
      x = nn.Dense(
          embedding_width, kernel_init=kernel_init, dtype=self.dtype
      )(x)
      x = nn.gelu(x)
      x = nn.LayerNorm(dtype=self.dtype)(x)
      x = jnp.einsum("ijk,lk->ijl", x, word_embeddings)
      x = x + self.param(
          "embedding_bias",
          nn.initializers.zeros,
          (1, 1, vocab_size),
          self.dtype,
      )
      return x
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the MaskedLanguageModelHead obj's underlying member method code with example usage: obj.__call__(x, masked_lm_positions, word_embeddings) and returns: x