class_docstr: Stores a representation of and computes the SSKernel function.
  
  K_L(A^dt, B^dt, C) corresponding to a discretized state space,
  where A is Normal + Low Rank (NPLR)
class_docstr_tok: ["Stores","a","representation","of","and","computes","the","SSKernel","function",".","K_L","(","A^dt",",","B^dt",",","C",")","corresponding","to","a","discretized","state","space",",","where","A","is","Normal","+","Low","Rank","(","NPLR",")"]
docstr: Forward pass.
  
  state: (B, H, N) initial state
  rate: sampling rate factor
  L: target length
  
  returns:
  (C, H, L) convolution kernel (generally C=1)
  (B, H, L) output from initial state
docstr_tok: ["Forward","pass",".","state",":","(","B",",","H",",","N",")","initial","state","rate",":","sampling","rate","factor","L",":","target","length","returns",":","(","C",",","H",",","L",")","convolution","kernel","(","generally","C=1",")","(","B",",","H",",","L",")","output","from","initial","state"]
code: def forward(self, state=None, rate=1.0, L=None):
      """Forward pass.
  
      state: (B, H, N) initial state
      rate: sampling rate factor
      L: target length
  
      returns:
      (C, H, L) convolution kernel (generally C=1)
      (B, H, L) output from initial state
      """
      # Initialize C~
      # if necessary (done in forward pass so it's on the correct device)
      if (
          self.L.item() == 0
          and self.l_max is not None
          and self.l_max > 0
      ):
          self._setup_C(self.l_max)
  
      # Handle sampling rate logic
      # The idea is that this kernel's length (in continuous units) is self.L,
      # while we are asked
      # to provide a kernel of length L at (relative) frequency rate
      if L is None:
          L = round(self.L.item() / rate)
  
      # Increase the internal length if needed
      continuous_L = round(rate * L)
      while continuous_L > self.L.item():
          self._setup_C(continuous_L)
      discrete_L = round(self.L.item() / rate)
  
      dt = torch.exp(self.log_dt) * rate
      B = _r2c(self.B)
      C = _r2c(self.C)
      P = _r2c(self.P)
      Q = P.conj()
      w = self._w()  # (n_ssm, N)
  
      # Address bandlimiting
      if self.bandlimit is not None:
          freqs = w.imag.abs() / (2 * math.pi)  # (H, N)
          freqs = dt[:, None] / rate * freqs  # (H, N)
          mask = torch.where(freqs < self.bandlimit * 0.5, 1, 0)
          C = C * mask
  
      # Get FFT nodes of right length
      omega, z = self._omega(
          discrete_L,
          dtype=w.dtype,
          device=w.device,
          cache=(rate == 1.0),
      )
  
      # Broadcast parameters to same hidden features H
      B = repeat(B, "1 t n -> 1 (v t) n", v=self.broadcast)
      P = repeat(P, "r t n -> r (v t) n", v=self.broadcast)
      Q = repeat(Q, "r t n -> r (v t) n", v=self.broadcast)
      w = repeat(w, "t n -> (v t) n", v=self.broadcast)
  
      # Augment B
      if state is not None:
          # Have to "unbilinear" the state to put it into the same "type" as B
          # Compute 1/dt * (I + dt/2 A) @ state
  
          # Can do this without expanding
          # (maybe minor speedup using conj symmetry in theory),
          # but it's easier to read this way
          s = (
              _conj(state) if state.size(-1) == self.N else state
          )  # (B H N)
          sA = s * _conj(w) - contract(  # (B H N)
              "bhm, rhm, rhn -> bhn", s, _conj(Q), _conj(P)
          )
          s = s / dt.unsqueeze(-1) + sA / 2
          s = s[..., : self.N]
  
          B = torch.cat([s, B], dim=-3)  # (B+1, H, N)
  
      # Incorporate dt into A
      w = w * dt.unsqueeze(-1)  # (H N)
  
      # Stack B and p, C and q for convenient batching
      B = torch.cat([B, P], dim=-3)  # (B+1+R, H, N)
      C = torch.cat([C, Q], dim=-3)  # (C+R, H, N)
  
      # Incorporate B and C batch dimensions
      v = B.unsqueeze(-3) * C.unsqueeze(-4)  # (B+1+R, C+R, H, N)
  
      # Calculate resolvent at omega
      if (
          has_cauchy_extension
          and z.dtype == torch.cfloat
          and not self.keops
      ):
          r = cauchy_mult(v, z, w, symmetric=True)
      elif has_pykeops:
          r = cauchy_conj(v, z, w)
      else:
          r = cauchy_naive(v, z, w)
      r = r * dt[None, None, :, None]  # (B+1+R, C+R, H, L)
  
      # Low-rank Woodbury correction
      if self.rank == 1:
          k_f = r[:-1, :-1, :, :] - r[:-1, -1:, :, :] * r[
              -1:, :-1, :, :
          ] / (1 + r[-1:, -1:, :, :])
      elif self.rank == 2:
          r00 = r[: -self.rank, : -self.rank, :, :]
          r01 = r[: -self.rank, -self.rank :, :, :]
          r10 = r[-self.rank :, : -self.rank, :, :]
          r11 = r[-self.rank :, -self.rank :, :, :]
          det = (1 + r11[:1, :1, :, :]) * (1 + r11[1:, 1:, :, :]) - r11[
              :1, 1:, :, :
          ] * r11[1:, :1, :, :]
          s = (
              r01[:, :1, :, :]
              * (1 + r11[1:, 1:, :, :])
              * r10[:1, :, :, :]
              + r01[:, 1:, :, :]
              * (1 + r11[:1, :1, :, :])
              * r10[1:, :, :, :]
              - r01[:, :1, :, :]
              * (r11[:1, 1:, :, :])
              * r10[1:, :, :, :]
              - r01[:, 1:, :, :]
              * (r11[1:, :1, :, :])
              * r10[:1, :, :, :]
          )
          s = s / det
          k_f = r00 - s
      else:
          r00 = r[: -self.rank, : -self.rank, :, :]
          r01 = r[: -self.rank, -self.rank :, :, :]
          r10 = r[-self.rank :, : -self.rank, :, :]
          r11 = r[-self.rank :, -self.rank :, :, :]
          r11 = rearrange(r11, "a b h n -> h n a b")
          r11 = torch.linalg.inv(
              torch.eye(self.rank, device=r.device) + r11
          )
          r11 = rearrange(r11, "h n a b -> a b h n")
          k_f = r00 - torch.einsum(
              "i j h n, j k h n, k l h n -> i l h n", r01, r11, r10
          )
  
      # Final correction for the bilinear transform
      k_f = k_f * 2 / (1 + omega)
  
      # Move from frequency to coefficients
      k = torch.fft.irfft(k_f, n=discrete_L)  # (B+1, C, H, L)
  
      # # Truncate to target length
      k = k[..., :L]
  
      if state is not None:
          k_state = k[:-1, :, :, :]  # (B, C, H, L)
      else:
          k_state = None
      k_B = k[-1, :, :, :]  # (C H L)
  
      return k_B, k_state
  
imports: import logging
  import math
  import os
  from functools import wraps
  from typing import Any, Callable, Optional
  import numpy
  import opt_einsum
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from espnet2.asr.state_spaces.components import Activation, DropoutNd, LinearActivation
label_desc: Use image node_id 6 for calling the SSKernelNPLR obj's underlying member method code with example usage: obj.forward(state, rate, L) and returns: k_B,  k_state