class_docstr: Attention pooling w/ latent query
      
code: def forward(self, x):
      B, N, C = x.shape
  
      if self.pos_embed is not None:
          # FIXME interpolate
          x = x + self.pos_embed.unsqueeze(0).to(x.dtype)
  
      q_latent = self.latent.expand(B, -1, -1)
      q = (
          self.q(q_latent)
          .reshape(B, self.latent_len, self.num_heads, self.head_dim)
          .transpose(1, 2)
      )
  
      kv = (
          self.kv(x)
          .reshape(B, N, 2, self.num_heads, self.head_dim)
          .permute(2, 0, 3, 1, 4)
      )
      k, v = kv.unbind(0)
  
      q, k = self.q_norm(q), self.k_norm(k)
  
      if self.fused_attn:
          x = F.scaled_dot_product_attention(q, k, v)
      else:
          q = q * self.scale
          attn = q @ k.transpose(-2, -1)
          attn = attn.softmax(dim=-1)
          x = attn @ v
      x = x.transpose(1, 2).reshape(B, self.latent_len, C)
      x = self.proj(x)
      x = self.proj_drop(x)
  
      x = x + self.mlp(self.norm(x))
  
      # optional pool if latent seq_len > 1 and pooled output is desired
      if self.pool == "token":
          x = x[:, 0]
      elif self.pool == "avg":
          x = x.mean(1)
      return x
  
file_path: pytorch-image-models/timm/layers/attention_pool.py
imports: from typing import Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from .config import use_fused_attn
  from .mlp import Mlp
  from .weight_init import trunc_normal_tf_
label_desc: Use image node_id 3 for calling the AttentionPoolLatent obj's underlying member method code with example usage: obj.forward(x) and returns: x