class_docstr: LSA: self attention within a group
      
code: def forward(self, x, size: Size_):
      # There are two implementations for this function, zero padding or mask. We don't observe obvious difference for
      # both. You can choose any one, we recommend forward_padding because it's neat. However,
      # the masking implementation is more reasonable and accurate.
      B, N, C = x.shape
      H, W = size
      x = x.view(B, H, W, C)
      pad_l = pad_t = 0
      pad_r = (self.ws - W % self.ws) % self.ws
      pad_b = (self.ws - H % self.ws) % self.ws
      x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
      _, Hp, Wp, _ = x.shape
      _h, _w = Hp // self.ws, Wp // self.ws
      x = x.reshape(B, _h, self.ws, _w, self.ws, C).transpose(2, 3)
      qkv = (
          self.qkv(x)
          .reshape(
              B,
              _h * _w,
              self.ws * self.ws,
              3,
              self.num_heads,
              C // self.num_heads,
          )
          .permute(3, 0, 1, 4, 2, 5)
      )
      q, k, v = qkv.unbind(0)
  
      if self.fused_attn:
          x = F.scaled_dot_product_attention(
              q,
              k,
              v,
              dropout_p=self.attn_drop.p if self.training else 0.0,
          )
      else:
          q = q * self.scale
          attn = q @ k.transpose(-2, -1)
          attn = attn.softmax(dim=-1)
          attn = self.attn_drop(attn)
          x = attn @ v
  
      x = x.transpose(2, 3).reshape(B, _h, _w, self.ws, self.ws, C)
      x = x.transpose(2, 3).reshape(B, _h * self.ws, _w * self.ws, C)
      if pad_r > 0 or pad_b > 0:
          x = x[:, :H, :W, :].contiguous()
      x = x.reshape(B, N, C)
      x = self.proj(x)
      x = self.proj_drop(x)
      return x
  
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 2 for calling the LocallyGroupedAttn obj's underlying member method code with example usage: obj.forward(x, size) and returns: x