class_docstr: Version using (complex) diagonal state matrix (S4D).
class_docstr_tok: ["Version","using","(","complex",")","diagonal","state","matrix","(","S4D",")","."]
code: def _setup_step(self):
      # These methods are organized
      # like this to be compatible with the NPLR kernel interface
      dt = torch.exp(self.log_dt)  # (H)
      B = _r2c(self.B)  # (H N)
      C = _r2c(self.C)  # (C H N)
      self.dC = C
      A = self._A()  # (H N)
  
      # Incorporate dt into A
      dtA = A * dt.unsqueeze(-1)  # (H N)
      if self.disc == "zoh":
          self.dA = torch.exp(dtA)  # (H N)
          self.dB = B * (torch.exp(dtA) - 1.0) / A  # (C H N)
      elif self.disc == "bilinear":
          self.dA = (1.0 + dtA / 2) / (1.0 - dtA / 2)
          self.dB = B * (1.0 - dtA / 2).reciprocal() * dt.unsqueeze(-1)
  
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 5 for calling the SSKernelDiag obj's underlying member method code with example usage: obj._setup_step() without return types