class_docstr: A simple softmax attention  operation
  LoFTR implementation reference: https://github.com/zju3dv/LoFTR/blob/2122156015b61fbb650e28b58a958e4d632b1058/src/loftr/loftr_module/linear_attention.py
class_docstr_tok: ["A","simple","softmax","attention","operation","LoFTR","implementation","reference",":","https",":","\/\/github.com\/zju3dv\/LoFTR\/blob\/2122156015b61fbb650e28b58a958e4d632b1058\/src\/loftr\/loftr_module\/linear_attention.py"]
docstr: Computes classical softmax full-attention between all queries and keys.
  
  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: [N, S1, H, D]
docstr_tok: ["Computes","classical","softmax","full-attention","between","all","queries","and","keys",".","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",":","[","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:
      """
      Computes classical softmax full-attention between all queries and keys.
  
      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: [N, S1, H, D]
      """
  
      scale_factor = 1.0 / queries.shape[3] ** 0.5  # irsqrt(D) scaling
      queries = queries * scale_factor
  
      qk = torch.einsum("NLHD, NSHD -> NLSH", queries, keys)
      if kv_mask is not None and q_mask is not None:
          qk.masked_fill_(
              ~(q_mask[:, :, None, None] * kv_mask[:, None, :, None]),
              float("-inf"),
          )
  
      attention = torch.softmax(qk, dim=2)
      attention = self.dropout(attention)
  
      queried_values = torch.einsum(
          "NLSH, NSHD -> NLHD", attention, values
      )
      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 SoftmaxAttention obj's underlying member method code with example usage: obj.forward(queries, keys, values, q_mask, kv_mask) and returns: queried_values