class_docstr: LSA: self attention within a group
      
code: def __init__(
      self, dim, num_heads=8, attn_drop=0.0, proj_drop=0.0, ws=1
  ):
      assert ws != 1
      super(LocallyGroupedAttn, self).__init__()
      assert (
          dim % num_heads == 0
      ), f"dim {dim} should be divided by num_heads {num_heads}."
  
      self.dim = dim
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5
      self.fused_attn = use_fused_attn()
  
      self.qkv = nn.Linear(dim, dim * 3, bias=True)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
      self.ws = ws
  
imports: import math
  from functools import partial
  from typing import Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import Mlp, DropPath, to_2tuple, trunc_normal_, use_fused_attn
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_module
  from ._registry import register_model, generate_default_cfgs
  from .vision_transformer import Attention
label_desc: Use image node_id 1 to create a new LocallyGroupedAttn object from inherited base classes: nn with example: obj = LocallyGroupedAttn(dim, num_heads, attn_drop, proj_drop, ws)