class_docstr: Version using (complex) diagonal state matrix (S4D).
class_docstr_tok: ["Version","using","(","complex",")","diagonal","state","matrix","(","S4D",")","."]
code: def __init__(
      self,
      A,
      B,
      C,
      log_dt,
      L=None,
      disc="bilinear",
      real_type="exp",
      lr=None,
      bandlimit=None,
  ):
      super().__init__()
      self.L = L
      self.disc = disc
      self.bandlimit = bandlimit
      self.real_type = real_type
  
      # Rank of low-rank correction
      assert A.size(-1) == C.size(-1)
      self.H = log_dt.size(-1)
      self.N = A.size(-1)
      assert A.size(-2) == B.size(
          -2
      )  # Number of independent SSMs trained
      assert self.H % A.size(-2) == 0
      self.n_ssm = A.size(-2)
      self.repeat = self.H // A.size(0)
  
      self.channels = C.shape[0]
      self.C = nn.Parameter(_c2r(_resolve_conj(C)))
  
      # Register parameters
      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("A", _c2r(A), lr_dict.get("A", lr))
      self.register("B", _c2r(B), lr_dict.get("B", lr))
      self.register(
          "inv_A_real", self._A_init(A.real), lr_dict.get("A", lr)
      )
      self.register("A_imag", A.imag, lr_dict.get("A", lr))
  
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 SSKernelDiag object from inherited base classes: OptimModule with example: obj = SSKernelDiag(A, B, C, log_dt, L, disc, real_type, lr, bandlimit)