class_docstr: Local state allows to have attention based only on data (no positional embedding),
  but while setting a constraint on the time window (e.g. decaying penalty term).
  Also a failed experiments with trying to provide some frequency based attention.
class_docstr_tok: ["Local","state","allows","to","have","attention","based","only","on","data","(","no","positional","embedding",")",",","but","while","setting","a","constraint","on","the","time","window","(","e.g",".","decaying","penalty","term",")",".","Also","a","failed","experiments","with","trying","to","provide","some","frequency","based","attention","."]
docstr: LocalState forward call
  
  Args:
      x (torch.Tensor): input tensor for LocalState
  
  Returns:
      Tensor
          Output after being run through LocalState layer.
docstr_tok: ["LocalState","forward","call","Args",":","x","(","torch.Tensor",")",":","input","tensor","for","LocalState","Returns",":","Tensor","Output","after","being","run","through","LocalState","layer","."]
code: def forward(self, x: torch.Tensor) -> torch.Tensor:
      r"""LocalState forward call
  
      Args:
          x (torch.Tensor): input tensor for LocalState
  
      Returns:
          Tensor
              Output after being run through LocalState layer.
      """
      B, C, T = x.shape
      heads = self.heads
      indexes = torch.arange(T, device=x.device, dtype=x.dtype)
      # left index are keys, right index are queries
      delta = indexes[:, None] - indexes[None, :]
  
      queries = self.query(x).view(B, heads, -1, T)
      keys = self.key(x).view(B, heads, -1, T)
      # t are keys, s are queries
      dots = torch.einsum("bhct,bhcs->bhts", keys, queries)
      dots /= math.sqrt(keys.shape[2])
      if self.ndecay:
          decays = torch.arange(
              1, self.ndecay + 1, device=x.device, dtype=x.dtype
          )
          decay_q = self.query_decay(x).view(B, heads, -1, T)
          decay_q = torch.sigmoid(decay_q) / 2
          decay_kernel = (
              -decays.view(-1, 1, 1)
              * delta.abs()
              / math.sqrt(self.ndecay)
          )
          dots += torch.einsum("fts,bhfs->bhts", decay_kernel, decay_q)
  
      # Kill self reference.
      dots.masked_fill_(
          torch.eye(T, device=dots.device, dtype=torch.bool), -100
      )
      weights = torch.softmax(dots, dim=2)
  
      content = self.content(x).view(B, heads, -1, T)
      result = torch.einsum("bhts,bhct->bhcs", weights, content)
      result = result.reshape(B, -1, T)
      return x + self.proj(result)
  
imports: import math
  import typing
  from typing import Any, Dict, List, Optional
  import torch
  from torch import nn
  from torch.nn import functional
label_desc: Use image node_id 2 for calling the _LocalState obj's underlying member method code with example usage: obj.forward(x) and returns: unknown