docstr: Initialize S4 module.
  
  d_state: the dimension of the state, also denoted by N
  l_max: the maximum kernel length, also denoted by L.
          Set l_max=None to always use a global kernel
  channels: can be interpreted as a number of "heads";
          the SSM is a map from a 1-dim to C-dim sequence.
          It's not recommended to change this unless desperate for things to tune;
          instead, increase d_model for larger models
  bidirectional: if True, convolution kernel will be two-sided
  
  Position-wise feedforward components:
  --------------------
  activation: activation in between SS and FF
  postact: activation after FF
  hyper_act: use a "hypernetwork" multiplication (experimental)
  dropout: standard dropout argument. tie_dropout=True ties the dropout
          mask across the sequence length, emulating nn.Dropout1d
  
  Other arguments:
  --------------------
  transposed: choose backbone axis ordering of
          (B, L, H) (if False) or (B, H, L) (if True)
          [B=batch size, L=sequence length, H=hidden dimension]
  gate: add gated activation (GSS)
  bottleneck: reduce SSM dimension (GSS)
  
  See the class SSKernel for the kernel constructor which accepts kernel_args.
  Relevant options that are worth considering
  and tuning include "mode" + "measure", "dt_min", "dt_max", "lr"
  
  Other options are all experimental and should not need to be configured
docstr_tok: ["Initialize","S4","module",".","d_state",":","the","dimension","of","the","state",",","also","denoted","by","N","l_max",":","the","maximum","kernel","length",",","also","denoted","by","L.","Set","l_max=None","to","always","use","a","global","kernel","channels",":","can","be","interpreted","as","a","number","of","``","heads","''",";","the","SSM","is","a","map","from","a","1-dim","to","C-dim","sequence",".","It","'s","not","recommended","to","change","this","unless","desperate","for","things","to","tune",";","instead",",","increase","d_model","for","larger","models","bidirectional",":","if","True",",","convolution","kernel","will","be","two-sided","Position-wise","feedforward","components",":","--","--","--","--","--","--","--","--","--","--","activation",":","activation","in","between","SS","and","FF","postact",":","activation","after","FF","hyper_act",":","use","a","``","hypernetwork","''","multiplication","(","experimental",")","dropout",":","standard","dropout","argument",".","tie_dropout=True","ties","the","dropout","mask","across","the","sequence","length",",","emulating","nn.Dropout1d","Other","arguments",":","--","--","--","--","--","--","--","--","--","--","transposed",":","choose","backbone","axis","ordering","of","(","B",",","L",",","H",")","(","if","False",")","or","(","B",",","H",",","L",")","(","if","True",")","[","B=batch","size",",","L=sequence","length",",","H=hidden","dimension","]","gate",":","add","gated","activation","(","GSS",")","bottleneck",":","reduce","SSM","dimension","(","GSS",")","See","the","class","SSKernel","for","the","kernel","constructor","which","accepts","kernel_args",".","Relevant","options","that","are","worth","considering","and","tuning","include","``","mode","''","+","``","measure","''",",","``","dt_min","''",",","``","dt_max","''",",","``","lr","''","Other","options","are","all","experimental","and","should","not","need","to","be","configured"]
code: def __init__(
      self,
      d_model,
      d_state=64,
      l_max=None,
      channels=1,
      bidirectional=False,
      # Arguments for position-wise feedforward components
      activation="gelu",
      postact="glu",
      hyper_act=None,
      dropout=0.0,
      tie_dropout=False,
      bottleneck=None,
      gate=None,
      transposed=True,
      verbose=False,
      # SSM Kernel arguments
      **kernel_args,
  ):
      """Initialize S4 module.
  
      d_state: the dimension of the state, also denoted by N
      l_max: the maximum kernel length, also denoted by L.
              Set l_max=None to always use a global kernel
      channels: can be interpreted as a number of "heads";
              the SSM is a map from a 1-dim to C-dim sequence.
              It's not recommended to change this unless desperate for things to tune;
              instead, increase d_model for larger models
      bidirectional: if True, convolution kernel will be two-sided
  
      Position-wise feedforward components:
      --------------------
      activation: activation in between SS and FF
      postact: activation after FF
      hyper_act: use a "hypernetwork" multiplication (experimental)
      dropout: standard dropout argument. tie_dropout=True ties the dropout
              mask across the sequence length, emulating nn.Dropout1d
  
      Other arguments:
      --------------------
      transposed: choose backbone axis ordering of
              (B, L, H) (if False) or (B, H, L) (if True)
              [B=batch size, L=sequence length, H=hidden dimension]
      gate: add gated activation (GSS)
      bottleneck: reduce SSM dimension (GSS)
  
      See the class SSKernel for the kernel constructor which accepts kernel_args.
      Relevant options that are worth considering
      and tuning include "mode" + "measure", "dt_min", "dt_max", "lr"
  
      Other options are all experimental and should not need to be configured
      """
      super().__init__()
      if verbose:
          log.info(
              f"Constructing S4 (H, N, L) = ({d_model}, {d_state}, {l_max})"
          )
  
      self.d_model = d_model
      self.H = d_model
      self.N = d_state
      self.L = l_max
      self.bidirectional = bidirectional
      self.channels = channels
      self.transposed = transposed
  
      self.gate = gate
      self.bottleneck = bottleneck
  
      if bottleneck is not None:
          self.H = self.H // bottleneck
          self.input_linear = LinearActivation(
              self.d_model,
              self.H,
              transposed=not self.transposed,
              activation=activation,
              activate=True,
          )
  
      if gate is not None:
          self.input_gate = LinearActivation(
              self.d_model,
              self.d_model * gate,
              transposed=not self.transposed,
              activation=activation,
              activate=True,
          )
          self.output_gate = LinearActivation(
              self.d_model * gate,
              self.d_model,
              transposed=self.transposed,
              activation=None,
              activate=False,
          )
  
      # optional multiplicative modulation GLU-style
      # https://arxiv.org/abs/2002.05202
      self.hyper = hyper_act is not None
      if self.hyper:
          channels *= 2
          self.hyper_activation = Activation(hyper_act)
  
      self.D = nn.Parameter(torch.randn(channels, self.H))
  
      if self.bidirectional:
          channels *= 2
  
      # SSM Kernel
      self.kernel = SSKernel(
          self.H,
          N=self.N,
          L=self.L,
          channels=channels,
          verbose=verbose,
          **kernel_args,
      )
  
      # Pointwise
      self.activation = Activation(activation)
      dropout_fn = DropoutNd if tie_dropout else nn.Dropout
      self.dropout = (
          dropout_fn(dropout) if dropout > 0.0 else nn.Identity()
      )
      # position-wise output transform to mix features
      self.output_linear = LinearActivation(
          self.H * self.channels,
          self.d_model * (1 if self.gate is None else self.gate),
          transposed=self.transposed,
          activation=postact,
          activate=True,
      )
  
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 S4 object from inherited base classes: nn with example: obj = S4(d_model, d_state, l_max, channels, bidirectional, activation, postact, hyper_act, dropout, tie_dropout, bottleneck, gate, transposed, verbose)