class_docstr: Lightweight multi-scale linear attention.
class_docstr_tok: ["Lightweight","multi-scale","linear","attention","."]
code: def relu_linear_att(self, qkv: torch.Tensor) -> torch.Tensor:
      B, _, H, W = list(qkv.size())
  
      if qkv.dtype == torch.float16:
          qkv = qkv.float()
  
      qkv = torch.reshape(qkv, (B, -1, 3 * self.dim, H * W))
      qkv = torch.transpose(qkv, -1, -2)
      q, k, v = (
          qkv[..., 0 : self.dim],
          qkv[..., self.dim : 2 * self.dim],
          qkv[..., 2 * self.dim :],
      )
  
      # lightweight linear attention
      q = self.kernel_func(q)
      k = self.kernel_func(k)
  
      # linear matmul
      trans_k = k.transpose(-1, -2)
  
      v = F.pad(v, (0, 1), mode="constant", value=1)
      kv = torch.matmul(trans_k, v)
      out = torch.matmul(q, kv)
      out = out[..., :-1] / (out[..., -1:] + self.eps)
  
      out = torch.transpose(out, -1, -2)
      out = torch.reshape(out, (B, -1, H, W))
      return out
  
imports: from __future__ import annotations
  import torch
  import torch.nn.functional
  from torch import nn
  from torch.cuda.amp import autocast
  from kornia.contrib.models.efficient_vit.nn.act import build_act
  from kornia.contrib.models.efficient_vit.nn.norm import build_norm
  from kornia.contrib.models.efficient_vit.utils import get_same_padding, val2tuple
label_desc: Use image node_id 2 for calling the LiteMLA obj's underlying member method code with example usage: obj.relu_linear_att(qkv) and returns: out