class_docstr: Version using (complex) diagonal state matrix (S4D).
class_docstr_tok: ["Version","using","(","complex",")","diagonal","state","matrix","(","S4D",")","."]
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, L, state=None, rate=1.0, u=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
      """
      dt = torch.exp(self.log_dt) * rate  # (H)
      C = _r2c(self.C)  # (C H N)
      A = self._A()  # (H N)
  
      B = _r2c(self.B)
      B = repeat(B, "t n -> 1 (v t) n", v=self.repeat)
  
      if self.bandlimit is not None:
          freqs = (
              dt[:, None] / rate * A.imag.abs() / (2 * math.pi)
          )  # (H, N)
          mask = torch.where(freqs < self.bandlimit * 0.5, 1, 0)
          C = C * mask
  
      # Incorporate dt into A
      A = repeat(A, "t n -> (v t) n", v=self.repeat)
      dtA = A * dt.unsqueeze(-1)  # (H N)
  
      # Augment B with state
      if state is not None:
          s = state / dt.unsqueeze(-1)
          if self.disc == "bilinear":
              s = s * (1.0 + dtA / 2)
          elif self.disc == "zoh":
              s = s * dtA * dtA.exp() / (dtA.exp() - 1.0)
          B = torch.cat([s, B], dim=-3)  # (1+B H N)
  
      C = (B[:, None, :, :] * C).view(-1, self.H, self.N)
      if self.disc == "zoh":
          # Power up
          C = C * (torch.exp(dtA) - 1.0) / A
          K = log_vandermonde(C, dtA, L)  # (H L)
      elif self.disc == "bilinear":
          C = (
              C * (1.0 - dtA / 2).reciprocal() * dt.unsqueeze(-1)
          )  # or * dtA / A
          dA = (1.0 + dtA / 2) / (1.0 - dtA / 2)
          K = log_vandermonde(C, dA.log(), L)
      elif self.disc == "dss":
          # Implementation from DSS meant for case
          # when real eigenvalues can be positive
          P = dtA.unsqueeze(-1) * torch.arange(
              L, device=C.device
          )  # [H N L]
          A_gt_0 = A.real > 0  # [N]
          if A_gt_0.any():
              with torch.no_grad():
                  P_max = dtA * (A_gt_0 * (L - 1))  # [H N]
              P = P - P_max.unsqueeze(-1)  # [H N L]
          S = P.exp()  # [H N L]
  
          dtA_neg = dtA * (1 - 2 * A_gt_0)  # [H N]
          num = dtA_neg.exp() - 1  # [H N]
          den = (dtA_neg * L).exp() - 1  # [H N]
  
          # Inline reciprocal function for DSS logic
          x = den * A
          x_conj = _resolve_conj(x)
          r = x_conj / (x * x_conj + 1e-7)
  
          C = C * num * r  # [C H N]
          K = contract("chn,hnl->chl", C, S).float()
      else:
          assert False, f"{self.disc} not supported"
  
      K = K.view(-1, self.channels, self.H, L)  # (1+B C H L)
      if state is not None:
          K_state = K[:-1, :, :, :]  # (B C H L)
      else:
          K_state = None
      K = K[-1, :, :, :]  # (C H L)
      return K, 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 4 for calling the SSKernelDiag obj's underlying member method code with example usage: obj.forward(L, state, rate, u) and returns: K,  K_state