code: def __init__(
      self,
      dim,
      mixer_cls=None,
      mlp_cls=None,
      norm_cls=nn.LayerNorm,
      dropout_cls=nn.Dropout,
      prenorm=True,
      resid_dropout=0.0,
      drop_path=0.0,
      fused_dropout_add_ln=False,
  ):
      super().__init__()
      self.prenorm = prenorm
      self.fused_dropout_add_ln = fused_dropout_add_ln
      if mixer_cls is None:
          mixer_cls = partial(MHA, num_heads=dim // 64)
      if mlp_cls is None:
          mlp_cls = partial(Mlp, hidden_features=4 * dim)
      self.mixer = mixer_cls(dim)
      self.dropout1 = dropout_cls(resid_dropout)
      self.drop_path1 = StochasticDepth(drop_path, mode="row")
      self.norm1 = norm_cls(dim)
      self.mlp = mlp_cls(dim)
      if not isinstance(self.mlp, nn.Identity):
          self.dropout2 = dropout_cls(resid_dropout)
          self.drop_path2 = StochasticDepth(drop_path, mode="row")
          self.norm2 = norm_cls(dim)
  
      if self.fused_dropout_add_ln:
          assert (
              dropout_add_layer_norm is not None
          ), "dropout_add_ln is not installed"
          assert isinstance(self.norm1, nn.LayerNorm) and isinstance(
              self.dropout1, nn.Dropout
          )
  
file_path: zoo/src/models/modules/block.py
imports: from typing import Optional
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torchvision.ops import StochasticDepth
  from src.models.modules.mha import MHA
  from src.models.modules.mlp import Mlp
label_desc: Use image node_id 1 to create a new Block object from inherited base classes: nn with example: obj = Block(dim, mixer_cls, mlp_cls, norm_cls, dropout_cls, prenorm, resid_dropout, drop_path, fused_dropout_add_ln)