class_docstr: Head used for classification with BERT.
  
  Attributes:
    num_outputs: Number of output classes.
    hidden_sizes: Size of hidden units in additional projections in the head.
    kernel_init: Kernel initialization.
    bias_init: Bias initialization.
    nonlinearity: Nonlinearity, ReLU by default.
    dtype: Model dtype.
class_docstr_tok: ["Head","used","for","classification","with","BERT",".","Attributes",":","num_outputs",":","Number","of","output","classes",".","hidden_sizes",":","Size","of","hidden","units","in","additional","projections","in","the","head",".","kernel_init",":","Kernel","initialization",".","bias_init",":","Bias","initialization",".","nonlinearity",":","Nonlinearity",",","ReLU","by","default",".","dtype",":","Model","dtype","."]
code: def __call__(self, x: jnp.ndarray, *, train: bool) -> jnp.ndarray:
      # Get CLS token:
      x = x[:, 0]
      hidden_sizes = self.hidden_sizes
      if isinstance(hidden_sizes, int):
          hidden_sizes = [hidden_sizes]
      for num_hid in hidden_sizes:
          # These intermediate layers are only used in BERT.
          x = nn.Dense(
              num_hid,
              kernel_init=bert_truncated_normal_initializer(),
              bias_init=self.bias_init,
          )(x)
          x = self.nonlinearity(x)
  
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      x = nn.Dense(
          self.num_outputs,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          name="output_projection",
      )(x)
      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 ClassificationHead obj's underlying member method code with example usage: obj.__call__(x) and returns: x