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: Set up dA, dB, dC discretized parameters for stepping.
docstr_tok: ["Set","up","dA",",","dB",",","dC","discretized","parameters","for","stepping","."]
code: def _setup_step(self, mode="dense"):
      """Set up dA, dB, dC discretized parameters for stepping."""
      self.dA, self.dB = self._setup_state()
  
      # Calculate original C
      C = _conj(_r2c(self.C))  # (H C N)
      if self.L.item() == 0:
          dC = C
      else:
          # self.C represents C_tilde
          dA_L = power(self.L.item(), self.dA)
          E = torch.eye(self.dA.size(-1)).to(dA_L)
  
          dC = torch.linalg.solve(
              E - dA_L.transpose(-1, -2),
              C.unsqueeze(-1),
          ).squeeze(-1)
      self.dC = dC
  
      # Do special preprocessing for different step modes
  
      self._step_mode = mode
      if mode == "linear":
          # Linear case: special step function for the state, we need to handle output
          # use conjugate symmetry by default, which affects the output projection
          self.dC = 2 * self.dC[:, :, : self.N]
      elif mode == "diagonal":
          # Eigendecomposition of the A matrix
          L, V = torch.linalg.eig(self.dA)
          V_inv = torch.linalg.inv(V)
          # Check that the eigendedecomposition is correct
          if self.verbose:
              print(
                  "Diagonalization error:",
                  torch.dist(V @ torch.diag_embed(L) @ V_inv, self.dA),
              )
  
          # Change the parameterization to diagonalize
          self.dA = L
          self.dB = contract("h n m, h m -> h n", V_inv, self.dB)
          self.dC = contract("h n m, c h n -> c h m", V, self.dC)
  
      elif mode == "dense":
          pass
      else:
          raise NotImplementedError(
              "NPLR Kernel step mode must be {'dense' | 'linear' | 'diagonal'}"
          )
  
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 11 for calling the SSKernelNPLR obj's underlying member method code with example usage: obj._setup_step(mode) without return types