class_docstr: Halo Attention
  
  Paper: `Scaling Local Self-Attention for Parameter Efficient Visual Backbones`
      - https://arxiv.org/abs/2103.12731
  
  The internal dimensions of the attention module are controlled by the interaction of several arguments.
    * the output dimension of the module is specified by dim_out, which falls back to input dim if not set
    * the value (v) dimension is set to dim_out // num_heads, the v projection determines the output dim
    * the query and key (qk) dimensions are determined by
      * num_heads * dim_head if dim_head is not None
      * num_heads * (dim_out * attn_ratio // num_heads) if dim_head is None
    * as seen above, attn_ratio determines the ratio of q and k relative to the output if dim_head not used
  
  Args:
      dim (int): input dimension to the module
      dim_out (int): output dimension of the module, same as dim if not set
      feat_size (Tuple[int, int]): size of input feature_map (not used, for arg compat with bottle/lambda)
      stride: output stride of the module, query downscaled if > 1 (default: 1).
      num_heads: parallel attention heads (default: 8).
      dim_head: dimension of query and key heads, calculated from dim_out * attn_ratio // num_heads if not set
      block_size (int): size of blocks. (default: 8)
      halo_size (int): size of halo overlap. (default: 3)
      qk_ratio (float): ratio of q and k dimensions to output dimension when dim_head not set. (default: 1.0)
      qkv_bias (bool) : add bias to q, k, and v projections
      avg_down (bool): use average pool downsample instead of strided query blocks
      scale_pos_embed (bool): scale the position embedding as well as Q @ K
class_docstr_tok: ["Halo","Attention","Paper",":","`","Scaling","Local","Self-Attention","for","Parameter","Efficient","Visual","Backbones","`","-","https",":","\/\/arxiv.org\/abs\/2103.12731","The","internal","dimensions","of","the","attention","module","are","controlled","by","the","interaction","of","several","arguments",".","*","the","output","dimension","of","the","module","is","specified","by","dim_out",",","which","falls","back","to","input","dim","if","not","set","*","the","value","(","v",")","dimension","is","set","to","dim_out","\/\/","num_heads",",","the","v","projection","determines","the","output","dim","*","the","query","and","key","(","qk",")","dimensions","are","determined","by","*","num_heads","*","dim_head","if","dim_head","is","not","None","*","num_heads","*","(","dim_out","*","attn_ratio","\/\/","num_heads",")","if","dim_head","is","None","*","as","seen","above",",","attn_ratio","determines","the","ratio","of","q","and","k","relative","to","the","output","if","dim_head","not","used","Args",":","dim","(","int",")",":","input","dimension","to","the","module","dim_out","(","int",")",":","output","dimension","of","the","module",",","same","as","dim","if","not","set","feat_size","(","Tuple","[","int",",","int","]",")",":","size","of","input","feature_map","(","not","used",",","for","arg","compat","with","bottle\/lambda",")","stride",":","output","stride","of","the","module",",","query","downscaled","if",">","1","(","default",":","1",")",".","num_heads",":","parallel","attention","heads","(","default",":","8",")",".","dim_head",":","dimension","of","query","and","key","heads",",","calculated","from","dim_out","*","attn_ratio","\/\/","num_heads","if","not","set","block_size","(","int",")",":","size","of","blocks",".","(","default",":","8",")","halo_size","(","int",")",":","size","of","halo","overlap",".","(","default",":","3",")","qk_ratio","(","float",")",":","ratio","of","q","and","k","dimensions","to","output","dimension","when","dim_head","not","set",".","(","default",":","1.0",")","qkv_bias","(","bool",")",":","add","bias","to","q",",","k",",","and","v","projections","avg_down","(","bool",")",":","use","average","pool","downsample","instead","of","strided","query","blocks","scale_pos_embed","(","bool",")",":","scale","the","position","embedding","as","well","as","Q","@","K"]
code: def forward(self, x):
      B, C, H, W = x.shape
      _assert(H % self.block_size == 0, "")
      _assert(W % self.block_size == 0, "")
      num_h_blocks = H // self.block_size
      num_w_blocks = W // self.block_size
      num_blocks = num_h_blocks * num_w_blocks
  
      q = self.q(x)
      # unfold
      q = q.reshape(
          -1,
          self.dim_head_qk,
          num_h_blocks,
          self.block_size_ds,
          num_w_blocks,
          self.block_size_ds,
      ).permute(0, 1, 3, 5, 2, 4)
      # B, num_heads * dim_head * block_size ** 2, num_blocks
      q = q.reshape(
          B * self.num_heads, self.dim_head_qk, -1, num_blocks
      ).transpose(1, 3)
      # B * num_heads, num_blocks, block_size ** 2, dim_head
  
      kv = self.kv(x)
      # Generate overlapping windows for kv. This approach is good for GPU and CPU. However, unfold() is not
      # lowered for PyTorch XLA so it will be very slow. See code at bottom of file for XLA friendly approach.
      # FIXME figure out how to switch impl between this and conv2d if XLA being used.
      kv = F.pad(
          kv,
          [
              self.halo_size,
              self.halo_size,
              self.halo_size,
              self.halo_size,
          ],
      )
      kv = (
          kv.unfold(2, self.win_size, self.block_size)
          .unfold(3, self.win_size, self.block_size)
          .reshape(
              B * self.num_heads,
              self.dim_head_qk + self.dim_head_v,
              num_blocks,
              -1,
          )
          .permute(0, 2, 3, 1)
      )
      k, v = torch.split(
          kv, [self.dim_head_qk, self.dim_head_v], dim=-1
      )
      # B * num_heads, num_blocks, win_size ** 2, dim_head_qk or dim_head_v
  
      if self.scale_pos_embed:
          attn = (
              q @ k.transpose(-1, -2) + self.pos_embed(q)
          ) * self.scale
      else:
          attn = (
              q @ k.transpose(-1, -2)
          ) * self.scale + self.pos_embed(q)
      # B * num_heads, num_blocks, block_size ** 2, win_size ** 2
      attn = attn.softmax(dim=-1)
  
      out = (attn @ v).transpose(
          1, 3
      )  # B * num_heads, dim_head_v, block_size ** 2, num_blocks
      # fold
      out = out.reshape(
          -1,
          self.block_size_ds,
          self.block_size_ds,
          num_h_blocks,
          num_w_blocks,
      )
      out = (
          out.permute(0, 3, 1, 4, 2)
          .contiguous()
          .view(
              B,
              self.dim_out_v,
              H // self.block_stride,
              W // self.block_stride,
          )
      )
      # B, dim_out, H // block_stride, W // block_stride
      out = self.pool(out)
      return out
  
file_path: pytorch-image-models/timm/layers/halo_attn.py
imports: from typing import List
  import torch
  from torch import nn
  import torch.nn.functional
  from .helpers import make_divisible
  from .weight_init import trunc_normal_
  from .trace_utils import _assert
label_desc: Use image node_id 3 for calling the HaloAttn obj's underlying member method code with example usage: obj.forward(x) and returns: out