class_docstr: Defines a fully connected neural network.
  
  The model assumes the input data has shape
  [batch_size_per_device, *input_shape] where input_shape may be of arbitrary
  rank. The model flatten the input before applying a dense layer.
  
  Attributes:
    num_outputs: Number of output classes.
    hid_sizes: Size of hidden units in each layer.
    kernel_init: Kernel initialization.
    bias_init: Bias initialization.
    output_proj_zero_init: Whether to initialize the output projection with
      zeros.
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    dtype: Model dtype.
class_docstr_tok: ["Defines","a","fully","connected","neural","network",".","The","model","assumes","the","input","data","has","shape","[","batch_size_per_device",",","*","input_shape","]","where","input_shape","may","be","of","arbitrary","rank",".","The","model","flatten","the","input","before","applying","a","dense","layer",".","Attributes",":","num_outputs",":","Number","of","output","classes",".","hid_sizes",":","Size","of","hidden","units","in","each","layer",".","kernel_init",":","Kernel","initialization",".","bias_init",":","Bias","initialization",".","output_proj_zero_init",":","Whether","to","initialize","the","output","projection","with","zeros",".","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","dtype",":","Model","dtype","."]
code: def __call__(self, x: jnp.ndarray, *, train: bool) -> jnp.ndarray:
      del train
      hid_sizes = self.hid_sizes
      if isinstance(hid_sizes, int):
          hid_sizes = [hid_sizes]
  
      x = get_bottleneck_representation(x, self.classifier)
  
      for num_hid in hid_sizes:
          x = nn.Dense(
              num_hid,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
          )(x)
          x = nn.relu(x)
  
      # Head.
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
  
      if self.output_proj_zero_init:
          output_proj_kernel_init = nn.initializers.zeros
          output_proj_bias_init = nn.initializers.zeros
      else:
          output_proj_kernel_init = self.kernel_init
          output_proj_bias_init = self.bias_init
  
      x = nn.Dense(
          self.num_outputs,
          kernel_init=output_proj_kernel_init,
          bias_init=output_proj_bias_init,
          name="output_projection",
      )(x)
      return x
  
imports: import functools
  from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.polyvit import polyvit_base_model
  from scenic.projects.vivit import model
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