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: Create parameters that allow fast linear stepping of state.
docstr_tok: ["Create","parameters","that","allow","fast","linear","stepping","of","state","."]
code: def _setup_linear(self):
      """Create parameters that allow fast linear stepping of state."""
      w = self._w()
      B = _r2c(self.B)  # (H N)
      P = _r2c(self.P)
      Q = P.conj()
  
      # Repeat w shape properly
      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)
  
      # Prepare Linear stepping
      dt = torch.exp(self.log_dt)
      D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
      R = (
          torch.eye(self.rank, dtype=w.dtype, device=w.device)
          + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
      )  # (H R R)
      Q_D = rearrange(Q * D, "r h n -> h r n")
      try:
          R = torch.linalg.solve(R, Q_D)  # (H R N)
      except:  # noqa
          R = torch.tensor(
              np.linalg.solve(
                  R.to(Q_D).contiguous().detach().cpu(),
                  Q_D.contiguous().detach().cpu(),
              )
          ).to(Q_D)
      R = rearrange(R, "h r n -> r h n")
  
      self.step_params = {
          "D": D,  # (H N)
          "R": R,  # (R H N)
          "P": P,  # (R H N)
          "Q": Q,  # (R H N)
          "B": B,  # (1 H N)
          "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
      }
  
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 7 for calling the SSKernelNPLR obj's underlying member method code with example usage: obj._setup_linear() without return types