class_docstr: Wrapper around SSKernel parameterizations.
  
  The SSKernel is expected to support the interface
  forward()
  default_state()
  _setup_step()
  step()
class_docstr_tok: ["Wrapper","around","SSKernel","parameterizations",".","The","SSKernel","is","expected","to","support","the","interface","forward","(",")","default_state","(",")","_setup_step","(",")","step","(",")"]
docstr: State Space Kernel which computes the convolution kernel $\\bar{K}$.
  
  H: Number of independent SSM copies;
      controls the size of the model. Also called d_model in the config.
  N: State size (dimensionality of parameters A, B, C).
      Also called d_state in the config.
      Generally shouldn't need to be adjusted and doens't affect speed much.
  L: Maximum length of convolution kernel, if known.
      Should work in the majority of cases even if not known.
  measure: Options for initialization of (A, B).
      For NPLR mode, recommendations are "legs",
      "fout", "hippo" (combination of both).
      For Diag mode, recommendations are "diag-inv",
      "diag-lin", "diag-legs", and "diag" (combination of diag-inv and diag-lin)
  rank: Rank of low-rank correction for NPLR mode.
      Needs to be increased for measure "legt"
  channels: C channels turns the SSM from a 1-dim to C-dim map;
      can think of it having C separate "heads" per SSM.
      This was partly a feature to make it easier to implement bidirectionality;
      it is recommended to set channels=1
      and adjust H to control parameters instead
  dt_min, dt_max: min and max values for the step size dt (\Delta)
  mode: Which kernel algorithm to use. 'nplr' is the full S4 model;
      'diag' is the simpler S4D; 'slow' is a dense version for testing
  n_ssm: Number of independent trainable (A, B) SSMs,
      e.g. n_ssm=1 means all A/B parameters are tied across
      the H different instantiations of C.
      n_ssm=None means all H SSMs are completely independent.
      Generally, changing this option can save parameters but doesn't affect
      performance or speed much. This parameter must divide H
  lr: Passing in a number (e.g. 0.001) sets
      attributes of SSM parameers (A, B, dt).
      A custom optimizer hook is needed to configure the optimizer
      to set the learning rates appropriately for these parameters.
docstr_tok: ["State","Space","Kernel","which","computes","the","convolution","kernel","$","\\\\bar","{","K","}","$",".","H",":","Number","of","independent","SSM","copies",";","controls","the","size","of","the","model",".","Also","called","d_model","in","the","config",".","N",":","State","size","(","dimensionality","of","parameters","A",",","B",",","C",")",".","Also","called","d_state","in","the","config",".","Generally","should","n't","need","to","be","adjusted","and","doens't","affect","speed","much",".","L",":","Maximum","length","of","convolution","kernel",",","if","known",".","Should","work","in","the","majority","of","cases","even","if","not","known",".","measure",":","Options","for","initialization","of","(","A",",","B",")",".","For","NPLR","mode",",","recommendations","are","``","legs","''",",","``","fout","''",",","``","hippo","''","(","combination","of","both",")",".","For","Diag","mode",",","recommendations","are","``","diag-inv","''",",","``","diag-lin","''",",","``","diag-legs","''",",","and","``","diag","''","(","combination","of","diag-inv","and","diag-lin",")","rank",":","Rank","of","low-rank","correction","for","NPLR","mode",".","Needs","to","be","increased","for","measure","``","legt","''","channels",":","C","channels","turns","the","SSM","from","a","1-dim","to","C-dim","map",";","can","think","of","it","having","C","separate","``","heads","''","per","SSM",".","This","was","partly","a","feature","to","make","it","easier","to","implement","bidirectionality",";","it","is","recommended","to","set","channels=1","and","adjust","H","to","control","parameters","instead","dt_min",",","dt_max",":","min","and","max","values","for","the","step","size","dt","(","\\Delta",")","mode",":","Which","kernel","algorithm","to","use",".","'nplr","'","is","the","full","S4","model",";","'diag","'","is","the","simpler","S4D",";","'slow","'","is","a","dense","version","for","testing","n_ssm",":","Number","of","independent","trainable","(","A",",","B",")","SSMs",",","e.g",".","n_ssm=1","means","all","A\/B","parameters","are","tied","across","the","H","different","instantiations","of","C.","n_ssm=None","means","all","H","SSMs","are","completely","independent",".","Generally",",","changing","this","option","can","save","parameters","but","does","n't","affect","performance","or","speed","much",".","This","parameter","must","divide","H","lr",":","Passing","in","a","number","(","e.g",".","0.001",")","sets","attributes","of","SSM","parameers","(","A",",","B",",","dt",")",".","A","custom","optimizer","hook","is","needed","to","configure","the","optimizer","to","set","the","learning","rates","appropriately","for","these","parameters","."]
code: def __init__(
      self,
      H,
      N=64,
      L=None,
      measure="legs",
      rank=1,
      channels=1,
      dt_min=0.001,
      dt_max=0.1,
      deterministic=False,
      lr=None,
      mode="nplr",
      n_ssm=None,
      verbose=False,
      measure_args={},
      **kernel_args,
  ):
      r"""State Space Kernel which computes the convolution kernel $\\bar{K}$.
  
      H: Number of independent SSM copies;
          controls the size of the model. Also called d_model in the config.
      N: State size (dimensionality of parameters A, B, C).
          Also called d_state in the config.
          Generally shouldn't need to be adjusted and doens't affect speed much.
      L: Maximum length of convolution kernel, if known.
          Should work in the majority of cases even if not known.
      measure: Options for initialization of (A, B).
          For NPLR mode, recommendations are "legs",
          "fout", "hippo" (combination of both).
          For Diag mode, recommendations are "diag-inv",
          "diag-lin", "diag-legs", and "diag" (combination of diag-inv and diag-lin)
      rank: Rank of low-rank correction for NPLR mode.
          Needs to be increased for measure "legt"
      channels: C channels turns the SSM from a 1-dim to C-dim map;
          can think of it having C separate "heads" per SSM.
          This was partly a feature to make it easier to implement bidirectionality;
          it is recommended to set channels=1
          and adjust H to control parameters instead
      dt_min, dt_max: min and max values for the step size dt (\Delta)
      mode: Which kernel algorithm to use. 'nplr' is the full S4 model;
          'diag' is the simpler S4D; 'slow' is a dense version for testing
      n_ssm: Number of independent trainable (A, B) SSMs,
          e.g. n_ssm=1 means all A/B parameters are tied across
          the H different instantiations of C.
          n_ssm=None means all H SSMs are completely independent.
          Generally, changing this option can save parameters but doesn't affect
          performance or speed much. This parameter must divide H
      lr: Passing in a number (e.g. 0.001) sets
          attributes of SSM parameers (A, B, dt).
          A custom optimizer hook is needed to configure the optimizer
          to set the learning rates appropriately for these parameters.
      """
      super().__init__()
      self.N = N
      self.H = H
      dtype, cdtype = torch.float, torch.cfloat
      self.channels = channels
      self.n_ssm = n_ssm if n_ssm is not None else H
      self.mode = mode
      self.verbose = verbose
      self.kernel_args = kernel_args
  
      # Generate dt
      if deterministic:
          log_dt = torch.exp(
              torch.linspace(math.log(dt_min), math.log(dt_max), H)
          )
      else:
          log_dt = torch.rand(self.H, dtype=dtype) * (
              math.log(dt_max) - math.log(dt_min)
          ) + math.log(dt_min)
  
      # Compute the preprocessed representation
      w, P, B, V = combination(
          measure, self.N, rank, self.n_ssm, **measure_args
      )
  
      # Broadcast C to have H channels
      if deterministic:
          C = torch.zeros(channels, self.H, self.N, dtype=cdtype)
          C[:, :, :1] = 1.0
          C = contract(
              "hmn, chn -> chm", V.conj().transpose(-1, -2), C
          )  # V^* C
      else:
          C = torch.randn(channels, self.H, self.N // 2, dtype=cdtype)
  
      # Broadcast other parameters to have n_ssm copies
      assert (
          self.n_ssm % B.size(-2) == 0
          and self.n_ssm % P.size(-2) == 0
          and self.n_ssm % w.size(-2) == 0
      )
      # Broadcast tensors to n_ssm copies
      # These will be the parameters,
      # so make sure tensors are materialized and contiguous
      B = (
          repeat(B, "t n -> (v t) n", v=self.n_ssm // B.size(-2))
          .clone()
          .contiguous()
      )
      P = (
          repeat(P, "r t n -> r (v t) n", v=self.n_ssm // P.size(-2))
          .clone()
          .contiguous()
      )
      w = (
          repeat(w, "t n -> (v t) n", v=self.n_ssm // w.size(-2))
          .clone()
          .contiguous()
      )
      C = C.contiguous()
  
      if mode == "nplr":
          self.kernel = SSKernelNPLR(
              w,
              P,
              B,
              C,
              log_dt,
              L=L,
              lr=lr,
              verbose=verbose,
              **kernel_args,
          )
      elif mode == "diag":
          if not measure.startswith("diag"):
              log.warning(
                  "Diagonal kernel (S4D) activated but initialization is not "
                  "intended for S4D. Set `measure` to 'diag-lin', 'diag-inv', or "
                  "'diag-legs' for the main variants, or 'diag' "
                  "for a combination of S4D-Lin and S4D-Inv."
              )
          C = C * repeat(B, "t n -> (v t) n", v=H // self.n_ssm)
          self.kernel = SSKernelDiag(
              w,
              B,
              C,
              log_dt,
              L=L,
              lr=lr,
              **kernel_args,
          )
      else:
          raise NotImplementedError(f"mode={mode} is not valid")
  
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 1 to create a new SSKernel object from inherited base classes: nn with example: obj = SSKernel(H, N, L, measure, rank, channels, dt_min, dt_max, deterministic, lr, mode, n_ssm, verbose, measure_args)