class_docstr: MLP.
class_docstr_tok: ["MLP","."]
code: def __call__(
      self,
      x: jnp.ndarray,  # Shape: (N, ..., X)
      *,
      train: bool = False,
      debug: bool = False,
  ) -> jnp.ndarray:  # Shape: (N, ..., E)
      is_batch_norm = self.normalization == "batch"
  
      y = x
  
      for _ in range(self.num_layers):
          y = nn.Dense(
              self.hidden_size,
              use_bias=not is_batch_norm,
              dtype=self.dtype,
          )(y)
  
          if self.normalization and self.normalization != "none":
              kwargs = {}
              if is_batch_norm:
                  kwargs["use_running_average"] = not train
  
              # We leave the default eps because the BatchNorm and LayerNorm are
              # always computed in FP32.
              y = normalizations.NORMALIZATIONS_BY_NAME[
                  self.normalization
              ](dtype=self.dtype)(
                  y, **kwargs
              )  # pytype: disable=wrong-keyword-args
  
          y = activations.ACTIVATIONS_BY_NAME[self.activation](y)
  
      y = nn.Dense(self.embedding_size, dtype=self.dtype)(y)
  
      # TODO(sacastro): dropout?
      # TODO(sacastro): norm at the end?
  
      if self.skip_connection:
          if x.shape[-1] != y.shape[-1]:
              x = nn.Dense(
                  self.embedding_size, use_bias=False, dtype=self.dtype
              )(x)
  
          return y + x
      else:
          return y
  
file_path: scenic/scenic/projects/lang4video/model/mlp.py
imports: from typing import Optional
  import flax.linen
  import jax.numpy
  from scenic.projects.lang4video.model import activations
  from scenic.projects.lang4video.model import normalizations
  from scenic.projects.lang4video.model.base_encoders import ImageEncoder
  from scenic.projects.lang4video.model.base_encoders import TextEncoder
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: unknown, y