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.
    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",".","dtype",":","Model","dtype","."]
docstr: Applies fully connected model on the input.
  
  Args:
    x: Input tensor.
    train: bool; Whether the model is running at train time.
    debug: bool; Whether the debug mode is enabled. debug=True enables model
      specific logging/storing some values using jax.host_callback.
  
  Returns:
    Unnormalized logits.
docstr_tok: ["Applies","fully","connected","model","on","the","input",".","Args",":","x",":","Input","tensor",".","train",":","bool",";","Whether","the","model","is","running","at","train","time",".","debug",":","bool",";","Whether","the","debug","mode","is","enabled",".","debug=True","enables","model","specific","logging\/storing","some","values","using","jax.host_callback",".","Returns",":","Unnormalized","logits","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ) -> jnp.ndarray:
      """Applies fully connected model on the input.
  
      Args:
        x: Input tensor.
        train: bool; Whether the model is running at train time.
        debug: bool; Whether the debug mode is enabled. debug=True enables model
          specific logging/storing some values using jax.host_callback.
  
      Returns:
        Unnormalized logits.
      """
      del train, debug
      hid_sizes = self.hid_sizes
      if isinstance(hid_sizes, int):
          hid_sizes = [hid_sizes]
      x = jnp.reshape(x, (x.shape[0], -1))
      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)
      x = nn.Dense(
          self.num_outputs,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          name="output_projection",
      )(x)
      return x
  
file_path: scenic/scenic/projects/baselines/fully_connected.py
imports: from typing import Callable, Iterable, Union, Sequence
  import flax.linen
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the FullyConnected obj's underlying member method code with example usage: obj.__call__(x) and returns: x