class_docstr: dim:
  num_heads:
class_docstr_tok: ["dim",":","num_heads",":"]
code: def __init__(
      self,
      dim,
      num_heads,
      mlp_ratio=4.0,
      qkv_bias=False,
      qk_scale=None,
      drop=0.0,
      attn_drop=0.0,
      drop_path=0.0,
      act_layer=nn.GELU,
      norm_layer=nn.LayerNorm,
      time_init="zeros",
      attention_style="frozen_in_time",
  ):
      super().__init__()
      self.attention_style = attention_style
      self.norm1 = norm_layer(dim)
      self.attn = VarAttention(
          dim,
          num_heads=num_heads,
          qkv_bias=qkv_bias,
          qk_scale=qk_scale,
          attn_drop=attn_drop,
          proj_drop=drop,
      )
      if self.attention_style != "bridge_former":
          self.timeattn = VarAttention(
              dim,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              qk_scale=qk_scale,
              attn_drop=attn_drop,
              proj_drop=drop,
              initialize=time_init,
          )
  
      # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
      self.drop_path = (
          DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
      )
      self.norm2 = norm_layer(dim)
      mlp_hidden_dim = int(dim * mlp_ratio)
      self.mlp = Mlp(
          in_features=dim,
          hidden_features=mlp_hidden_dim,
          act_layer=act_layer,
          drop=drop,
      )
      if self.attention_style != "bridge_former":
          self.norm3 = norm_layer(dim)
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  import logging
label_desc: Use image node_id 1 to create a new SpaceTimeBlock object from inherited base classes: nn with example: obj = SpaceTimeBlock(dim, num_heads, mlp_ratio, qkv_bias, qk_scale, drop, attn_drop, drop_path, act_layer, norm_layer, time_init, attention_style)