class_docstr: MLPMixer.
  
  Implements
  MLP-Mixer: An all-MLP Architecture for Vision
  https://arxiv.org/abs/2105.01601
code: def __init__(
      self,
      img_height: int,
      img_width: int,
      in_channels: int,
      patch_size: int = 16,
      embed_size: int = 512,
      token_size: int = 2048,
      channel_dim: int = 256,
      num_layers: int = 8,
      dropout: float = 0.0,
      avg_pool: bool = True,
  ):
      super().__init__()
      assert (img_height % patch_size == 0) and (
          img_width % patch_size == 0
      )
  
      self._input_shape = (in_channels, img_height, img_width)
      n_patches = int(img_height * img_width / (patch_size**2))
  
      self.patch_conv = nn.Conv2d(
          in_channels=in_channels,
          out_channels=embed_size,
          kernel_size=patch_size,
          stride=patch_size,
      )
  
      self.mixer_blocks = nn.ModuleList(
          [
              MixerBlock(
                  embed_size=embed_size,
                  n_patches=n_patches,
                  token_dim=token_size,
                  channel_dim=channel_dim,
                  dropout=dropout,
              )
              for _ in range(num_layers)
          ]
      )
  
      self.layer_norm = nn.LayerNorm(normalized_shape=embed_size)
  
      self.avg_pool = avg_pool
      if self.avg_pool:
          self._output_shape = torch.Size((embed_size,))
      else:
          self._output_shape = torch.Size((n_patches, embed_size))
  
imports: from typing import Tuple, Union
  import torch
  import torch.nn
  from ludwig.utils.torch_utils import LudwigModule
label_desc: Use image node_id 1 to create a new MLPMixer object from inherited base classes: LudwigModule with example: obj = MLPMixer(img_height, img_width, in_channels, patch_size, embed_size, token_size, channel_dim, num_layers, dropout, avg_pool)