class_docstr: Linear attention operation from: https://arxiv.org/pdf/2006.16236.pdf
  Canonical implementation reference: https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py
  LoFTR implementation reference: https://github.com/zju3dv/LoFTR/blob/2122156015b61fbb650e28b58a958e4d632b1058/src/loftr/loftr_module/linear_attention.py
class_docstr_tok: ["Linear","attention","operation","from",":","https",":","\/\/arxiv.org\/pdf\/2006.16236.pdf","Canonical","implementation","reference",":","https",":","\/\/github.com\/idiap\/fast-transformers\/blob\/master\/fast_transformers\/attention\/linear_attention.py","LoFTR","implementation","reference",":","https",":","\/\/github.com\/zju3dv\/LoFTR\/blob\/2122156015b61fbb650e28b58a958e4d632b1058\/src\/loftr\/loftr_module\/linear_attention.py"]
docstr: Args:
      queries (torch.Tensor): [N, S1, H, D]
      keys (torch.Tensor): [N, S2, H, D]
      values (torch.Tensor): [N, S2, H, D]
      q_mask (torch.Tensor): [N, S1] (optional)
      kv_mask (torch.Tensor): [N, S2] (optional)
  Returns:
      queried_values (torch.Tensor): [N, S1, H, D]
docstr_tok: ["Args",":","queries","(","torch.Tensor",")",":","[","N",",","S1",",","H",",","D","]","keys","(","torch.Tensor",")",":","[","N",",","S2",",","H",",","D","]","values","(","torch.Tensor",")",":","[","N",",","S2",",","H",",","D","]","q_mask","(","torch.Tensor",")",":","[","N",",","S1","]","(","optional",")","kv_mask","(","torch.Tensor",")",":","[","N",",","S2","]","(","optional",")","Returns",":","queried_values","(","torch.Tensor",")",":","[","N",",","S1",",","H",",","D","]"]
code: def forward(
      self,
      queries: Tensor,
      keys: Tensor,
      values: Tensor,
      q_mask: Optional[Tensor] = None,
      kv_mask: Optional[Tensor] = None,
  ) -> Tensor:
      """
      Args:
          queries (torch.Tensor): [N, S1, H, D]
          keys (torch.Tensor): [N, S2, H, D]
          values (torch.Tensor): [N, S2, H, D]
          q_mask (torch.Tensor): [N, S1] (optional)
          kv_mask (torch.Tensor): [N, S2] (optional)
      Returns:
          queried_values (torch.Tensor): [N, S1, H, D]
      """
      queries = self.feature_map_fn(queries)
      keys = self.feature_map_fn(keys)
  
      if q_mask is not None:
          queries = queries * q_mask[:, :, None, None]
      if kv_mask is not None:
          keys = keys * kv_mask[:, :, None, None]
          values = values * kv_mask[:, :, None, None]
  
      # mitigates fp16 overflows
      values_length = values.shape[1]
      values = values / values_length
      kv = torch.einsum("NSHD, NSHV -> NHDV", keys, values)
      z = 1 / (
          torch.einsum("NLHD, NHD -> NLH", queries, keys.sum(dim=1))
          + self.eps
      )
      # rescale at the end to account for fp16 mitigation
      queried_values = (
          torch.einsum("NLHD, NHDV, NLH -> NLHV", queries, kv, z)
          * values_length
      )
      return queried_values
  
imports: import math
  from functools import partial
  from typing import Callable, Dict, Iterable, List, Optional, Tuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torchvision.models.optical_flow.raft
  from torch import Tensor
  from torchvision.models._api import register_model, Weights, WeightsEnum
  from torchvision.models._utils import handle_legacy_interface
  from torchvision.models.optical_flow._utils import grid_sample, make_coords_grid, upsample_flow
  from torchvision.ops import Conv2dNormActivation
  from torchvision.prototype.transforms._presets import StereoMatching
label_desc: Use image node_id 2 for calling the LinearAttention obj's underlying member method code with example usage: obj.forward(queries, keys, values, q_mask, kv_mask) and returns: queried_values