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: Initialize kernel.
  
  L: Maximum length; this module computes an SSM kernel of length L
  A is represented by diag(w) - PP^*
  w: (S, N) diagonal part
  P: (R, S, N) low-rank part
  
  B: (S, N)
  C: (C, H, N)
  dt: (H) timescale per feature
  lr: [dict | float | None] hook to set lr of special parameters (A, B, dt)
  
  Dimensions:
  N (or d_state): state size
  H (or d_model): total SSM copies
  S (or n_ssm): number of trainable copies of (A, B, dt); must divide H
  R (or rank): rank of low-rank part
  C (or channels): system is 1-dim to C-dim
  
  The forward pass of this Module returns a tensor of shape (C, H, L)
  
  Note: tensor shape N here denotes half the true state size,
      because of conjugate symmetry
docstr_tok: ["Initialize","kernel",".","L",":","Maximum","length",";","this","module","computes","an","SSM","kernel","of","length","L","A","is","represented","by","diag","(","w",")","-","PP^","*","w",":","(","S",",","N",")","diagonal","part","P",":","(","R",",","S",",","N",")","low-rank","part","B",":","(","S",",","N",")","C",":","(","C",",","H",",","N",")","dt",":","(","H",")","timescale","per","feature","lr",":","[","dict","|","float","|","None","]","hook","to","set","lr","of","special","parameters","(","A",",","B",",","dt",")","Dimensions",":","N","(","or","d_state",")",":","state","size","H","(","or","d_model",")",":","total","SSM","copies","S","(","or","n_ssm",")",":","number","of","trainable","copies","of","(","A",",","B",",","dt",")",";","must","divide","H","R","(","or","rank",")",":","rank","of","low-rank","part","C","(","or","channels",")",":","system","is","1-dim","to","C-dim","The","forward","pass","of","this","Module","returns","a","tensor","of","shape","(","C",",","H",",","L",")","Note",":","tensor","shape","N","here","denotes","half","the","true","state","size",",","because","of","conjugate","symmetry"]
code: def __init__(
      self,
      w,
      P,
      B,
      C,
      log_dt,
      L=None,  # starting/maximum length of kernel
      lr=None,
      verbose=False,
      keops=False,
      real_type="exp",  # ['none' | 'exp' | 'relu' | sigmoid']
      real_tolerance=1e-3,
      bandlimit=None,
  ):
      """Initialize kernel.
  
      L: Maximum length; this module computes an SSM kernel of length L
      A is represented by diag(w) - PP^*
      w: (S, N) diagonal part
      P: (R, S, N) low-rank part
  
      B: (S, N)
      C: (C, H, N)
      dt: (H) timescale per feature
      lr: [dict | float | None] hook to set lr of special parameters (A, B, dt)
  
      Dimensions:
      N (or d_state): state size
      H (or d_model): total SSM copies
      S (or n_ssm): number of trainable copies of (A, B, dt); must divide H
      R (or rank): rank of low-rank part
      C (or channels): system is 1-dim to C-dim
  
      The forward pass of this Module returns a tensor of shape (C, H, L)
  
      Note: tensor shape N here denotes half the true state size,
          because of conjugate symmetry
      """
      super().__init__()
      self.verbose = verbose
      self.keops = keops
      self.bandlimit = bandlimit
      self.real_type = real_type
      self.real_tolerance = real_tolerance
  
      # Rank of low-rank correction
      self.rank = P.shape[-3]
      assert w.size(-1) == P.size(-1) == B.size(-1) == C.size(-1)
      self.H = log_dt.size(-1)
      self.N = w.size(-1)
  
      # Check different SSM inits
      assert w.size(-2) == P.size(-2) == B.size(-2)  # n_ssm
      assert self.H % w.size(0) == 0
      self.n_ssm = w.size(0)
      self.broadcast = self.H // w.size(
          0
      )  # Each trainable SSM needs to be duplicated this many times
  
      # Broadcast everything to correct shapes
      C = C.expand(
          torch.broadcast_shapes(C.shape, (1, self.H, self.N))
      )  # (C, H, N)
      B = B.unsqueeze(0)  # (1, 1, N)
  
      # Register parameters
      self.C = nn.Parameter(_c2r(_resolve_conj(C)))
      if lr is None or isinstance(lr, float):
          lr_dict = {}
      else:
          lr_dict, lr = lr, None
      self.register("log_dt", log_dt, lr_dict.get("dt", lr))
      self.register("B", _c2r(B), lr_dict.get("B", lr))
      self.register("P", _c2r(P), lr_dict.get("A", lr))
      self.register(
          "inv_w_real", self._w_init(w.real), lr_dict.get("A", lr)
      )
      self.register("w_imag", w.imag, lr_dict.get("A", lr))
  
      self.l_max = L
      self.register_buffer("L", torch.tensor(0))
  
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 3 to create a new SSKernelNPLR object from inherited base classes: OptimModule with example: obj = SSKernelNPLR(w, P, B, C, log_dt, L, lr, verbose, keops, real_type, real_tolerance, bandlimit)