docstr: Multi-Head linear attention proposed in "Transformers are RNNs"
  Args:
      queries: [N, L, H, D]
      keys: [N, S, H, D]
      values: [N, S, H, D]
      q_mask: [N, L]
      kv_mask: [N, S]
  Returns:
      queried_values: (N, L, H, D)
docstr_tok: ["Multi-Head","linear","attention","proposed","in","``","Transformers","are","RNNs","''","Args",":","queries",":","[","N",",","L",",","H",",","D","]","keys",":","[","N",",","S",",","H",",","D","]","values",":","[","N",",","S",",","H",",","D","]","q_mask",":","[","N",",","L","]","kv_mask",":","[","N",",","S","]","Returns",":","queried_values",":","(","N",",","L",",","H",",","D",")"]
code: def forward(
      self,
      queries: Tensor,
      keys: Tensor,
      values: Tensor,
      q_mask: Optional[Tensor] = None,
      kv_mask: Optional[Tensor] = None,
  ) -> Tensor:
      """Multi-Head linear attention proposed in "Transformers are RNNs"
      Args:
          queries: [N, L, H, D]
          keys: [N, S, H, D]
          values: [N, S, H, D]
          q_mask: [N, L]
          kv_mask: [N, S]
      Returns:
          queried_values: (N, L, H, D)
      """
      Q = self.feature_map(queries)
      K = self.feature_map(keys)
  
      # set padded position to zero
      if q_mask is not None:
          Q = Q * q_mask[:, :, None, None]
      if kv_mask is not None:
          K = K * kv_mask[:, :, None, None]
          values = values * kv_mask[:, :, None, None]
  
      v_length = values.size(1)
      values = values / v_length  # prevent fp16 overflow
      KV = torch.einsum("nshd,nshv->nhdv", K, values)  # (S,D)' @ S,V
      Z = 1 / (
          torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1)) + self.eps
      )
      queried_values = (
          torch.einsum("nlhd,nhdv,nlh->nlhv", Q, KV, Z) * v_length
      )
  
      return queried_values.contiguous()
  
file_path: kornia/kornia/feature/loftr/loftr_module/linear_attention.py
imports: from typing import Optional
  import torch
  from torch.nn import Dropout
  from kornia.core import Module, Tensor
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