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: Step one time step as a recurrent model.
  
  Version of the step function that has time O(N) instead of O(N^2) per step,
  which takes advantage of the DPLR form and bilinear discretization.
  
  Unfortunately, as currently implemented it's about 2x slower
  because it calls several sequential operations.
  Perhaps a fused CUDA kernel implementation would be much faster
  
  u: (H) input
  state: (H, N/2) state with conjugate pairs
    Optionally, the state can have last dimension N
  Returns: same shape as state
docstr_tok: ["Step","one","time","step","as","a","recurrent","model",".","Version","of","the","step","function","that","has","time","O","(","N",")","instead","of","O","(","N^2",")","per","step",",","which","takes","advantage","of","the","DPLR","form","and","bilinear","discretization",".","Unfortunately",",","as","currently","implemented","it","'s","about","2x","slower","because","it","calls","several","sequential","operations",".","Perhaps","a","fused","CUDA","kernel","implementation","would","be","much","faster","u",":","(","H",")","input","state",":","(","H",",","N\/2",")","state","with","conjugate","pairs","Optionally",",","the","state","can","have","last","dimension","N","Returns",":","same","shape","as","state"]
code: def _step_state_linear(self, u=None, state=None):
      """Step one time step as a recurrent model.
  
      Version of the step function that has time O(N) instead of O(N^2) per step,
      which takes advantage of the DPLR form and bilinear discretization.
  
      Unfortunately, as currently implemented it's about 2x slower
      because it calls several sequential operations.
      Perhaps a fused CUDA kernel implementation would be much faster
  
      u: (H) input
      state: (H, N/2) state with conjugate pairs
        Optionally, the state can have last dimension N
      Returns: same shape as state
      """
      C = _r2c(self.C)  # View used for dtype/device
  
      if u is None:  # Special case used to find dA
          u = torch.zeros(self.H, dtype=C.dtype, device=C.device)
      if state is None:  # Special case used to find dB
          state = torch.zeros(
              self.H, self.N, dtype=C.dtype, device=C.device
          )
  
      step_params = self.step_params.copy()
      if (
          state.size(-1) == self.N
      ):  # Only store half of the conjugate pairs; should be true by default
          # There should be a slightly faster way using conjugate symmetry
          def contract_fn(p, x, y):
              return contract(
                  "r h n, r h m, ... h m -> ... h n",
                  _conj(p),
                  _conj(x),
                  _conj(y),
              )[
                  ..., : self.N
              ]  # inner outer product
  
      else:
          assert state.size(-1) == 2 * self.N
          step_params = {k: _conj(v) for k, v in step_params.items()}
          # Worth setting up a contract_expression in default_state
          # if we want to use this at inference time for stepping
  
          def contract_fn(p, x, y):
              return contract(
                  "r h n, r h m, ... h m -> ... h n", p, x, y
              )  # inner outer product
  
      D = step_params["D"]  # (H N)
      E = step_params["E"]  # (H N)
      R = step_params["R"]  # (R H N)
      P = step_params["P"]  # (R H N)
      Q = step_params["Q"]  # (R H N)
      B = step_params["B"]  # (1 H N)
  
      new_state = E * state - contract_fn(P, Q, state)  # (B H N)
      new_state = new_state + 2.0 * B * u.unsqueeze(-1)  # (B H N)
      new_state = D * (new_state - contract_fn(P, R, new_state))
  
      return new_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 8 for calling the SSKernelNPLR obj's underlying member method code with example usage: obj._step_state_linear(u, state) and returns: new_state, contract, contract