class_docstr: Defines a simple convolutional neural network.
  
  The model assumes the input shape is [batch, H, W, C].
  
  Attributes:
    num_outputs: Number of output classes.
    output_projection_type: Type of the output projection layer.
    num_filters: Number of filters in each layer.
    kernel_sizes: Size of kernel in each layer.
    use_bias: If add bias in each layer.
    kernel_init: Kernel initialization.
    bias_init: Bias initialization.
    dtype: Model JAX dtype.
class_docstr_tok: ["Defines","a","simple","convolutional","neural","network",".","The","model","assumes","the","input","shape","is","[","batch",",","H",",","W",",","C","]",".","Attributes",":","num_outputs",":","Number","of","output","classes",".","output_projection_type",":","Type","of","the","output","projection","layer",".","num_filters",":","Number","of","filters","in","each","layer",".","kernel_sizes",":","Size","of","kernel","in","each","layer",".","use_bias",":","If","add","bias","in","each","layer",".","kernel_init",":","Kernel","initialization",".","bias_init",":","Bias","initialization",".","dtype",":","Model","JAX","dtype","."]
docstr: Applies SimpleCNN on the input data.
  
  Args:
    x: Input tensor.
    train: Unused.
    debug: Unused.
  
  Returns:
    Unnormalized logits.
docstr_tok: ["Applies","SimpleCNN","on","the","input","data",".","Args",":","x",":","Input","tensor",".","train",":","Unused",".","debug",":","Unused",".","Returns",":","Unnormalized","logits","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ) -> jnp.ndarray:
      """Applies SimpleCNN on the input data.
  
      Args:
        x: Input tensor.
        train: Unused.
        debug: Unused.
  
      Returns:
        Unnormalized logits.
      """
      del train, debug
      use_bias = True if self.use_bias is None else self.use_bias
      if not isinstance(use_bias, Iterable):
          use_bias = [use_bias] * len(self.num_filters)
      for n_filters, kernel_size, use_bias in zip(
          self.num_filters, self.kernel_sizes, use_bias
      ):
          x = nn.Conv(
              features=n_filters,
              kernel_size=(kernel_size, kernel_size),
              strides=(1, 1),
              use_bias=use_bias,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
              dtype=self.dtype,
          )(x)
          x = nn.relu(x)
  
      # Head
      if self.output_projection_type == "reduce_mean":
          x = jnp.mean(x, axis=(1, 2))
  
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      x = nn.Dense(
          self.num_outputs,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          dtype=self.dtype,
          name="output_projection",
      )(x)
      return x
  
file_path: scenic/scenic/projects/baselines/simple_cnn.py
imports: from typing import Iterable, Callable, Sequence, Optional, Union
  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.base_models.segmentation_model import SegmentationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the SimpleCNN obj's underlying member method code with example usage: obj.__call__(x) and returns: x