class_docstr: Simple MLP.
docstr: Forward module.
  
  Args:
    x: array in shape batch_size x num_tokens x hidden_dim
  Returns:
    batch_size x num_tokens x out_dim: 2 for softmax
docstr_tok: ["Forward","module",".","Args",":","x",":","array","in","shape","batch_size","x","num_tokens","x","hidden_dim","Returns",":","batch_size","x","num_tokens","x","out_dim",":","2","for","softmax"]
code: def __call__(self, x):
      """Forward module.
  
      Args:
        x: array in shape batch_size x num_tokens x hidden_dim
      Returns:
        batch_size x num_tokens x out_dim: 2 for softmax
      """
      x = nn.LayerNorm()(x)
      for i in range(self.num_layers):
          x = nn.Dense(self.hidden_dim, name=f"linear.{i}")(x)
          x = nn.gelu(x)
      return x
  
imports: import functools
  from typing import Any, Callable, Iterable, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.objectvivit.object_attention import ObjectBlock
label_desc: Use image node_id 1 for calling the MLP obj's underlying member method code with example usage: obj.__call__(x) and returns: x