self-contained-inference (#1)

- fix: make ConMamba-small-ca self-contained for inference (df1804fb060f9af7e0f673c3283e90629c8717e8)


Co-authored-by: ferran <guasch@users.noreply.huggingface.co>

conmambamamba_debug_catalan_small_1k_unigram_inference.yaml

@@ -1,386 +0,0 @@
-# ############################################################################
-# Model: E2E ASR with Transformer
-# Encoder: ConMamba Encoder
-# Decoder: Mamba Decoder + (CTC/ATT joint) beamsearch + TransformerLM
-# Tokens: unigram
-# losses: CTC + KLdiv (Label Smoothing loss)
-# Training: Librispeech 960h
-# Authors: Xilin Jiang
-# ############################################################################
-# Seed needs to be set at top of yaml, before objects with parameters are made
-
-seed: 7780
-__set_seed: !apply:torch.manual_seed [!ref <seed>]
-project: Mamba-ASR
-experiment: conmambamamba_S_S2S_catalan_1k_unigram
-output_folder: !ref /gpfs/projects/bsc88/speech/TTS/outputs/checkpoints/Mamba_ASR/catalan/<experiment>/<seed>
-output_wer_folder: !ref <output_folder>/
-save_folder: !ref <output_folder>/save
-train_log: !ref <output_folder>/train_log.txt
-
-# Language model (LM) pretraining
-# NB: To avoid mismatch, the speech recognizer must be trained with the same
-# tokenizer used for LM training. Here, we download everything from the
-# speechbrain HuggingFace repository. However, a local path pointing to a
-# directory containing the lm.ckpt and tokenizer.ckpt may also be specified
-# instead. E.g if you want to use your own LM / tokenizer.
-pretrained_lm_tokenizer_path: /gpfs/projects/bsc88/speech/TTS/repos/Mamba-ASR/catalan/lm
-
-# Data files
-# If RIRS_NOISES dir exists in /localscratch/xxx_corpus/RIRS_NOISES
-# then data_folder_rirs should be /localscratch/xxx_corpus
-# otherwise the dataset will automatically be downloaded
-# data_folder_rirs: !ref <data_folder>
-train_splits: ["train-clean-100"]
-dev_splits: ["dev-clean"]
-test_splits: ["test-clean"]
-skip_prep: True
-train_csv: /gpfs/projects/bsc88/speech/TTS/repos/Mamba-ASR/catalan/catalan_train_merged_duration.csv
-valid_csv: /gpfs/projects/bsc88/speech/TTS/repos/Mamba-ASR/catalan/catalan_dev_2h.csv
-test_csv:
-    - /gpfs/projects/bsc88/speech/TTS/repos/Mamba-ASR/catalan/catalan_test_merged_duration.csv
-
-skip_train: True
-no_lm: True
-
-# Training parameters
-# To make Transformers converge, the global bath size should be large enough.
-# The global batch size is computed as batch_size * n_gpus * grad_accumulation_factor.
-# Empirically, we found that this value should be >= 128.
-# Please, set your parameters accordingly.
-number_of_epochs: 110    # En original debemos ponerle minimo 110
-batch_size: 30 # This works for 2x GPUs with 32GB
-ctc_weight: 0.6 #0.3
-grad_accumulation_factor: 1 #1 valor normal
-max_grad_norm: 5.0
-loss_reduction: 'batchmean'
-sorting: random
-num_workers: 8 #8 old version
-precision: bf16 # bf16, fp16 or fp32
-avg_checkpoints: 10 # Number of checkpoints to average for evaluation
-
-# stages related parameters
-# stage_one_epochs: 90
-lr_adam: 0.001 #0.001
-# lr_sgd: 0.000025
-
-# Feature parameters
-sample_rate: 16000
-n_fft: 400
-n_mels: 80
-
-# This setup works well for V100 32GB GPU, adapts it to your needs.
-# Or turn it off (but training speed will decrease)
-dynamic_batching: False #True
-max_batch_length_train: 900 #900 old version
-max_batch_length_val: 100 #100 we reduce it as the beam is much wider (VRAM)
-num_bucket: 200
-shuffle: True # if true re-creates batches at each epoch shuffling examples.
-batch_ordering: random
-max_batch_ex: 256 #128 old version
-
-dynamic_batch_sampler_train:
-    max_batch_length: !ref <max_batch_length_train>
-    num_buckets: !ref <num_bucket>
-    shuffle: !ref <shuffle>
-    batch_ordering: !ref <batch_ordering>
-    max_batch_ex: !ref <max_batch_ex>
-
-dynamic_batch_sampler_valid:
-    max_batch_length: !ref <max_batch_length_val>
-    num_buckets: !ref <num_bucket>
-    shuffle: !ref <shuffle>
-    batch_ordering: !ref <batch_ordering>
-    max_batch_ex: !ref <max_batch_ex>
-
-# Dataloader options
-train_dataloader_opts:
-    batch_size: !ref <batch_size>
-    shuffle: True
-    num_workers: !ref <num_workers>
-
-valid_dataloader_opts:
-    batch_size: 1
-
-test_dataloader_opts:
-    batch_size: 1
-
-####################### Model parameters ###########################
-# Transformer
-d_model: 144
-nhead: 4
-num_encoder_layers: 12
-num_decoder_layers: 4
-d_ffn: 1024
-transformer_dropout: 0.1
-activation: !name:torch.nn.GELU
-output_neurons: 1000
-
-# Outputs
-blank_index: 0
-label_smoothing: 0.1 #0.0
-pad_index: 0
-bos_index: 1
-eos_index: 2
-
-# Decoding parameters
-min_decode_ratio: 0.0
-max_decode_ratio: 1.0
-valid_search_interval: 15 #Cambiar esto para ver el WER
-valid_beam_size: 10
-test_beam_size: 24 #66
-lm_weight: 0.30 #0.60
-ctc_weight_decode: 0.70 #0.40
-
-# Mamba parameters
-
-d_state: 16
-expand: 2
-d_conv: 4
-bidirectional: True
-mamba_config:
-    d_state: !ref <d_state>
-    expand: !ref <expand>
-    d_conv: !ref <d_conv>
-    bidirectional: !ref <bidirectional>
-
-############################## models ################################
-
-CNN: !new:speechbrain.lobes.models.convolution.ConvolutionFrontEnd
-    input_shape: (8, 10, 80)
-    num_blocks: 2
-    num_layers_per_block: 1
-    out_channels: (64, 32)
-    kernel_sizes: (3, 3)
-    strides: (2, 2)
-    residuals: (False, False)
-
-Transformer: !new:modules.TransformerASR.TransformerASR # yamllint disable-line rule:line-length
-    input_size: 640
-    tgt_vocab: !ref <output_neurons>
-    d_model: !ref <d_model>
-    nhead: !ref <nhead> # unused
-    num_encoder_layers: !ref <num_encoder_layers>
-    num_decoder_layers: !ref <num_decoder_layers>
-    d_ffn: !ref <d_ffn>
-    dropout: !ref <transformer_dropout>
-    activation: !ref <activation>
-    encoder_module: conmamba
-    decoder_module: mamba
-    attention_type: RelPosMHAXL # unused
-    normalize_before: True
-    causal: False
-    mamba_config: !ref <mamba_config>
-
-# This is the TransformerLM that is used according to the Huggingface repository
-# Visit the HuggingFace model corresponding to the pretrained_lm_tokenizer_path
-# For more details about the model!
-# NB: It has to match the pre-trained TransformerLM!!
-lm_model: !new:speechbrain.lobes.models.transformer.TransformerLM.TransformerLM # yamllint disable-line rule:line-length
-    vocab: !ref <output_neurons>
-    d_model: 768
-    nhead: 12
-    num_encoder_layers: 12
-    num_decoder_layers: 0
-    d_ffn: 3072
-    dropout: 0.0
-    activation: !name:torch.nn.GELU
-    normalize_before: False
-
-
-
-tokenizer: !new:sentencepiece.SentencePieceProcessor
-
-ctc_lin: !new:speechbrain.nnet.linear.Linear
-    input_size: !ref <d_model>
-    n_neurons: !ref <output_neurons>
-
-seq_lin: !new:speechbrain.nnet.linear.Linear
-    input_size: !ref <d_model>
-    n_neurons: !ref <output_neurons>
-
-normalize: !new:speechbrain.processing.features.InputNormalization
-    norm_type: global
-    update_until_epoch: 4
-
-modules:
-    CNN: !ref <CNN>
-    Transformer: !ref <Transformer>
-    seq_lin: !ref <seq_lin>
-    ctc_lin: !ref <ctc_lin>
-    normalize: !ref <normalize>
-
-model: !new:torch.nn.ModuleList
-    - [!ref <CNN>, !ref <Transformer>, !ref <seq_lin>, !ref <ctc_lin>]
-
-# define two optimizers here for two-stage training
-Adam: !name:torch.optim.Adam
-    lr: !ref <lr_adam>
-    betas: (0.9, 0.98)
-    eps: 0.000000001
-
-# Scorer
-ctc_scorer: !new:speechbrain.decoders.scorer.CTCScorer
-    eos_index: !ref <eos_index>
-    blank_index: !ref <blank_index>
-    ctc_fc: !ref <ctc_lin>
-
-
-transformerlm_scorer: !new:speechbrain.decoders.scorer.TransformerLMScorer
-    language_model: !ref <lm_model>
-    temperature: 1.15
-
-scorer_test_search: !new:speechbrain.decoders.scorer.ScorerBuilder
-    full_scorers: [!ref <transformerlm_scorer>, !ref <ctc_scorer>]
-    weights:
-        ctc: !ref <ctc_weight_decode>
-        transformerlm: !ref <lm_weight>
-
-scorer_valid_search: !new:speechbrain.decoders.scorer.ScorerBuilder
-    full_scorers: [!ref <ctc_scorer>]
-    weights:
-        ctc: !ref <ctc_weight_decode>
-
-valid_search: !new:speechbrain.decoders.S2STransformerBeamSearcher
-    modules: [!ref <Transformer>, !ref <seq_lin>]
-    bos_index: !ref <bos_index>
-    eos_index: !ref <eos_index>
-    min_decode_ratio: !ref <min_decode_ratio>
-    max_decode_ratio: !ref <max_decode_ratio>
-    beam_size: !ref <valid_beam_size>
-    using_eos_threshold: False
-    length_normalization: True
-    scorer: !ref <scorer_valid_search>
-
-test_search: !new:speechbrain.decoders.S2STransformerBeamSearcher
-    modules: [!ref <Transformer>, !ref <seq_lin>]
-    bos_index: !ref <bos_index>
-    eos_index: !ref <eos_index>
-    min_decode_ratio: !ref <min_decode_ratio>
-    max_decode_ratio: !ref <max_decode_ratio>
-    beam_size: !ref <test_beam_size>
-    temperature: 1.15
-    using_eos_threshold: False
-    length_normalization: True
-    scorer: !ref <scorer_test_search>
-
-
-log_softmax: !new:torch.nn.LogSoftmax
-    dim: -1
-
-ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
-    blank_index: !ref <blank_index>
-    reduction: !ref <loss_reduction>
-
-seq_cost: !name:speechbrain.nnet.losses.kldiv_loss
-    label_smoothing: !ref <label_smoothing>
-    reduction: !ref <loss_reduction>
-
-noam_annealing: !new:speechbrain.nnet.schedulers.NoamScheduler
-    lr_initial: !ref <lr_adam>
-    n_warmup_steps: !ref 25000 // <grad_accumulation_factor>
-
-checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
-    checkpoints_dir: !ref <save_folder>
-    recoverables:
-        model: !ref <model>
-        noam_scheduler: !ref <noam_annealing>
-        normalizer: !ref <normalize>
-        counter: !ref <epoch_counter>
-
-epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
-    limit: !ref <number_of_epochs>
-
-# Speed perturbation
-speed_changes: [95, 100, 105]  # List of speed changes for time-stretching
-
-speed_perturb: !new:speechbrain.augment.time_domain.SpeedPerturb
-    orig_freq: !ref <sample_rate>
-    speeds: !ref <speed_changes>
-
-# Time Drop
-time_drop_length_low: 15  # Min length for temporal chunk to drop in spectrogram
-time_drop_length_high: 25  # Max length for temporal chunk to drop in spectrogram
-time_drop_count_low: 2  #4 Min number of chunks to drop in time in the spectrogram
-time_drop_count_high: 2  #4 Max number of chunks to drop in time in the spectrogram
-time_drop_replace: "mean"  # Method of dropping chunks
-
-time_drop: !new:speechbrain.augment.freq_domain.SpectrogramDrop
-    drop_length_low: !ref <time_drop_length_low>
-    drop_length_high: !ref <time_drop_length_high>
-    drop_count_low: !ref <time_drop_count_low>
-    drop_count_high: !ref <time_drop_count_high>
-    replace: !ref <time_drop_replace>
-    dim: 1
-
-# Frequency Drop
-freq_drop_length_low: 10  # Min length for chunks to drop in frequency in the spectrogram
-freq_drop_length_high: 20  # Max length for chunks to drop in frequency in the spectrogram
-freq_drop_count_low: 2  #4 Min number of chunks to drop in frequency in the spectrogram
-freq_drop_count_high: 2  #4 Max number of chunks to drop in frequency in the spectrogram
-freq_drop_replace: "mean"  # Method of dropping chunks
-
-freq_drop: !new:speechbrain.augment.freq_domain.SpectrogramDrop
-    drop_length_low: !ref <freq_drop_length_low>
-    drop_length_high: !ref <freq_drop_length_high>
-    drop_count_low: !ref <freq_drop_count_low>
-    drop_count_high: !ref <freq_drop_count_high>
-    replace: !ref <freq_drop_replace>
-    dim: 2
-
-# Time warp
-time_warp_window: 5  # Length of time warping window
-time_warp_mode: "bicubic"  # Time warping method
-
-time_warp: !new:speechbrain.augment.freq_domain.Warping
-    warp_window: !ref <time_warp_window>
-    warp_mode: !ref <time_warp_mode>
-    dim: 1
-
-fea_augment: !new:speechbrain.augment.augmenter.Augmenter
-    parallel_augment: False
-    concat_original: False
-    repeat_augment: 1
-    shuffle_augmentations: False
-    min_augmentations: 3
-    max_augmentations: 3
-    augment_prob: 1.0
-    augmentations: [
-        !ref <time_drop>,
-        !ref <freq_drop>,
-        !ref <time_warp>]
-
-compute_features: !new:speechbrain.lobes.features.Fbank
-    sample_rate: !ref <sample_rate>
-    n_fft: !ref <n_fft>
-    n_mels: !ref <n_mels>
-
-train_logger: !new:speechbrain.utils.train_logger.FileTrainLogger
-    save_file: !ref <train_log>
-
-error_rate_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
-acc_computer: !name:speechbrain.utils.Accuracy.AccuracyStats
-
-# The pretrainer allows a mapping between pretrained files and instances that
-# are declared in the yaml. E.g here, we will download the file lm.ckpt
-# and it will be loaded into "lm" which is pointing to the <lm_model> defined
-# before.
-pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
-    collect_in: !ref <save_folder>
-    loadables:
-        lm: !ref <lm_model>
-        tokenizer: !ref <tokenizer>
-    paths:
-        lm: !ref <pretrained_lm_tokenizer_path>/lm.ckpt
-        tokenizer: !ref <pretrained_lm_tokenizer_path>/tokenizer.ckpt
-
-use_wandb: False #False
-resume: False
-wandb_logger: !name:speechbrain.utils.train_logger.WandBLogger
-    initializer: !name:wandb.init
-    entity: BSC
-    project: !ref <project>
-    name: !ref <experiment>
-    dir: !ref <save_folder>
-    reinit: true
-    resume: !ref <resume>

inference.yaml

@@ -0,0 +1,131 @@
+# ConMamba-small-ca standalone inference configuration.
+# Decoding uses the LM-free beam search validated in ASR Arena.
+
+seed: 7780
+__set_seed: !apply:torch.manual_seed [!ref <seed>]
+
+# The inference entry point overrides model_dir with the repository directory,
+# so these assets resolve correctly even when invoked from another directory.
+model_dir: .
+model_checkpoint: !ref <model_dir>/model.ckpt
+tokenizer_checkpoint: !ref <model_dir>/tokenizer.ckpt
+normalizer_checkpoint: !ref <model_dir>/normalizer.ckpt
+
+# Audio features
+sample_rate: 16000
+n_fft: 400
+n_mels: 80
+
+# Model architecture
+d_model: 144
+nhead: 4
+num_encoder_layers: 12
+num_decoder_layers: 4
+d_ffn: 1024
+transformer_dropout: 0.1
+activation: !name:torch.nn.GELU
+output_neurons: 1000
+
+# Token indices
+blank_index: 0
+pad_index: 0
+bos_index: 1
+eos_index: 2
+
+# LM-free decoding parameters
+min_decode_ratio: 0.0
+max_decode_ratio: 1.0
+beam_size: 10
+ctc_weight_decode: 0.70
+
+# Mamba parameters
+d_state: 16
+expand: 2
+d_conv: 4
+bidirectional: True
+mamba_config:
+    d_state: !ref <d_state>
+    expand: !ref <expand>
+    d_conv: !ref <d_conv>
+    bidirectional: !ref <bidirectional>
+
+# Acoustic model
+CNN: !new:speechbrain.lobes.models.convolution.ConvolutionFrontEnd
+    input_shape: (8, 10, 80)
+    num_blocks: 2
+    num_layers_per_block: 1
+    out_channels: (64, 32)
+    kernel_sizes: (3, 3)
+    strides: (2, 2)
+    residuals: (False, False)
+
+Transformer: !new:modules.TransformerASR.TransformerASR
+    input_size: 640
+    tgt_vocab: !ref <output_neurons>
+    d_model: !ref <d_model>
+    nhead: !ref <nhead>
+    num_encoder_layers: !ref <num_encoder_layers>
+    num_decoder_layers: !ref <num_decoder_layers>
+    d_ffn: !ref <d_ffn>
+    dropout: !ref <transformer_dropout>
+    activation: !ref <activation>
+    encoder_module: conmamba
+    decoder_module: mamba
+    attention_type: RelPosMHAXL
+    normalize_before: True
+    causal: False
+    mamba_config: !ref <mamba_config>
+
+tokenizer: !new:sentencepiece.SentencePieceProcessor
+
+ctc_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: !ref <output_neurons>
+
+seq_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: !ref <output_neurons>
+
+normalize: !new:speechbrain.processing.features.InputNormalization
+    norm_type: global
+    update_until_epoch: 4
+
+modules:
+    CNN: !ref <CNN>
+    Transformer: !ref <Transformer>
+    seq_lin: !ref <seq_lin>
+    ctc_lin: !ref <ctc_lin>
+    normalize: !ref <normalize>
+
+model: !new:torch.nn.ModuleList
+    - [!ref <CNN>, !ref <Transformer>, !ref <seq_lin>, !ref <ctc_lin>]
+
+# CTC/attention beam search, without an external language model
+ctc_scorer: !new:speechbrain.decoders.scorer.CTCScorer
+    eos_index: !ref <eos_index>
+    blank_index: !ref <blank_index>
+    ctc_fc: !ref <ctc_lin>
+
+search_scorer: !new:speechbrain.decoders.scorer.ScorerBuilder
+    full_scorers: [!ref <ctc_scorer>]
+    weights:
+        ctc: !ref <ctc_weight_decode>
+
+beam_search: !new:speechbrain.decoders.S2STransformerBeamSearcher
+    modules: [!ref <Transformer>, !ref <seq_lin>]
+    bos_index: !ref <bos_index>
+    eos_index: !ref <eos_index>
+    min_decode_ratio: !ref <min_decode_ratio>
+    max_decode_ratio: !ref <max_decode_ratio>
+    beam_size: !ref <beam_size>
+    using_eos_threshold: False
+    length_normalization: True
+    scorer: !ref <search_scorer>
+
+log_softmax: !new:torch.nn.LogSoftmax
+    dim: -1
+
+compute_features: !new:speechbrain.lobes.features.Fbank
+    sample_rate: !ref <sample_rate>
+    n_fft: !ref <n_fft>
+    n_mels: !ref <n_mels>

model.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be1a4250fe4f60d7bf5308bdc3f06430be7d195bff3f6dcbe8f483df1b4285b4
+oid sha256:cb075002ffd3258f9cb3ad7f127c55d8bcf61f5cc3ba5a28a7b9573d38f07b60
 size 52989622

modules/Conformer.py

@@ -0,0 +1,1094 @@
+"""Conformer implementation.
+
+Authors
+-------
+* Jianyuan Zhong 2020
+* Samuele Cornell 2021
+* Sylvain de Langen 2023
+"""
+
+import warnings
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import speechbrain as sb
+from speechbrain.nnet.activations import Swish
+from speechbrain.nnet.attention import (
+    MultiheadAttention,
+    PositionalwiseFeedForward,
+    RelPosMHAXL,
+)
+from speechbrain.nnet.hypermixing import HyperMixing
+from speechbrain.nnet.normalization import LayerNorm
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+
+
+@dataclass
+class ConformerEncoderLayerStreamingContext:
+    """Streaming metadata and state for a `ConformerEncoderLayer`.
+
+    The multi-head attention and Dynamic Chunk Convolution require to save some
+    left context that gets inserted as left padding.
+
+    See :class:`.ConvolutionModule` documentation for further details.
+    """
+
+    mha_left_context_size: int
+    """For this layer, specifies how many frames of inputs should be saved.
+    Usually, the same value is used across all layers, but this can be modified.
+    """
+
+    mha_left_context: Optional[torch.Tensor] = None
+    """Left context to insert at the left of the current chunk as inputs to the
+    multi-head attention. It can be `None` (if we're dealing with the first
+    chunk) or `<= mha_left_context_size` because for the first few chunks, not
+    enough left context may be available to pad.
+    """
+
+    dcconv_left_context: Optional[torch.Tensor] = None
+    """Left context to insert at the left of the convolution according to the
+    Dynamic Chunk Convolution method.
+
+    Unlike `mha_left_context`, here the amount of frames to keep is fixed and
+    inferred from the kernel size of the convolution module.
+    """
+
+
+@dataclass
+class ConformerEncoderStreamingContext:
+    """Streaming metadata and state for a `ConformerEncoder`."""
+
+    dynchunktrain_config: DynChunkTrainConfig
+    """Dynamic Chunk Training configuration holding chunk size and context size
+    information."""
+
+    layers: List[ConformerEncoderLayerStreamingContext]
+    """Streaming metadata and state for each layer of the encoder."""
+
+
+class ConvolutionModule(nn.Module):
+    """This is an implementation of convolution module in Conformer.
+
+    Arguments
+    ---------
+    input_size : int
+        The expected size of the input embedding dimension.
+    kernel_size: int, optional
+        Kernel size of non-bottleneck convolutional layer.
+    bias: bool, optional
+        Whether to use bias in the non-bottleneck conv layer.
+    activation: torch.nn.Module
+         Activation function used after non-bottleneck conv layer.
+    dropout: float, optional
+         Dropout rate.
+    causal: bool, optional
+         Whether the convolution should be causal or not.
+    dilation: int, optional
+         Dilation factor for the non bottleneck conv layer.
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> net = ConvolutionModule(512, 3)
+    >>> output = net(x)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        input_size,
+        kernel_size=31,
+        bias=True,
+        activation=Swish,
+        dropout=0.0,
+        causal=False,
+        dilation=1,
+    ):
+        super().__init__()
+
+        self.kernel_size = kernel_size
+        self.causal = causal
+        self.dilation = dilation
+
+        if self.causal:
+            self.padding = (kernel_size - 1) * 2 ** (dilation - 1)
+        else:
+            self.padding = (kernel_size - 1) * 2 ** (dilation - 1) // 2
+
+        self.layer_norm = nn.LayerNorm(input_size)
+        self.bottleneck = nn.Sequential(
+            # pointwise
+            nn.Conv1d(
+                input_size, 2 * input_size, kernel_size=1, stride=1, bias=bias
+            ),
+            nn.GLU(dim=1),
+        )
+        # depthwise
+        self.conv = nn.Conv1d(
+            input_size,
+            input_size,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=self.padding,
+            dilation=dilation,
+            groups=input_size,
+            bias=bias,
+        )
+
+        # BatchNorm in the original Conformer replaced with a LayerNorm due to
+        # https://github.com/speechbrain/speechbrain/pull/1329
+        # see discussion
+        # https://github.com/speechbrain/speechbrain/pull/933#issuecomment-1033367884
+
+        self.after_conv = nn.Sequential(
+            nn.LayerNorm(input_size),
+            activation(),
+            # pointwise
+            nn.Linear(input_size, input_size, bias=bias),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """Applies the convolution to an input tensor `x`.
+
+        Arguments
+        ---------
+        x: torch.Tensor
+            Input tensor to the convolution module.
+        mask: torch.Tensor, optional
+            Mask to be applied over the output of the convolution using
+            `masked_fill_`, if specified.
+        dynchunktrain_config: DynChunkTrainConfig, optional
+            If specified, makes the module support Dynamic Chunk Convolution
+            (DCConv) as implemented by
+            `Dynamic Chunk Convolution for Unified Streaming and Non-Streaming Conformer ASR <https://www.amazon.science/publications/dynamic-chunk-convolution-for-unified-streaming-and-non-streaming-conformer-asr>`_.
+            This allows masking future frames while preserving better accuracy
+            than a fully causal convolution, at a small speed cost.
+            This should only be used for training (or, if you know what you're
+            doing, for masked evaluation at inference time), as the forward
+            streaming function should be used at inference time.
+
+        Returns
+        -------
+        out: torch.Tensor
+            The output tensor.
+        """
+
+        if dynchunktrain_config is not None:
+            # chances are chunking+causal is unintended; i don't know where it
+            # may make sense, but if it does to you, feel free to implement it.
+            assert (
+                not self.causal
+            ), "Chunked convolution not supported with causal padding"
+
+            assert (
+                self.dilation == 1
+            ), "Current DynChunkTrain logic does not support dilation != 1"
+
+            # in a causal convolution, which is not the case here, an output
+            # frame would never be able to depend on a input frame from any
+            # point in the future.
+
+            # but with the dynamic chunk convolution, we instead use a "normal"
+            # convolution but where, for any output frame, the future beyond the
+            # "current" chunk gets masked.
+            # see the paper linked in the documentation for details.
+
+            chunk_size = dynchunktrain_config.chunk_size
+            batch_size = x.shape[0]
+
+            # determine the amount of padding we need to insert at the right of
+            # the last chunk so that all chunks end up with the same size.
+            if x.shape[1] % chunk_size != 0:
+                final_right_padding = chunk_size - (x.shape[1] % chunk_size)
+            else:
+                final_right_padding = 0
+
+            # -> [batch_size, t, in_channels]
+            out = self.layer_norm(x)
+
+            # -> [batch_size, in_channels, t] for the CNN
+            out = out.transpose(1, 2)
+
+            # -> [batch_size, in_channels, t] (pointwise)
+            out = self.bottleneck(out)
+
+            # -> [batch_size, in_channels, lc+t+final_right_padding]
+            out = F.pad(out, (self.padding, final_right_padding), value=0)
+
+            # now, make chunks with left context.
+            # as a recap to what the above padding and this unfold do, consider
+            # each a/b/c letter represents a frame as part of chunks a, b, c.
+            # consider a chunk size of 4 and a kernel size of 5 (padding=2):
+            #
+            # input seq: 00aaaabbbbcc00
+            # chunk #1:  00aaaa
+            # chunk #2:      aabbbb
+            # chunk #3:          bbcc00
+            #
+            # a few remarks here:
+            # - the left padding gets inserted early so that the unfold logic
+            #   works trivially
+            # - the right 0-padding got inserted as the number of time steps
+            #   could not be evenly split in `chunk_size` chunks
+
+            # -> [batch_size, in_channels, num_chunks, lc+chunk_size]
+            out = out.unfold(2, size=chunk_size + self.padding, step=chunk_size)
+
+            # as we manually disable padding in the convolution below, we insert
+            # right 0-padding to the chunks, e.g. reusing the above example:
+            #
+            # chunk #1:  00aaaa00
+            # chunk #2:      aabbbb00
+            # chunk #3:          bbcc0000
+
+            # -> [batch_size, in_channels, num_chunks, lc+chunk_size+rpad]
+            out = F.pad(out, (0, self.padding), value=0)
+
+            # the transpose+flatten effectively flattens chunks into the batch
+            # dimension to be processed into the time-wise convolution. the
+            # chunks will later on be unflattened.
+
+            # -> [batch_size, num_chunks, in_channels, lc+chunk_size+rpad]
+            out = out.transpose(1, 2)
+
+            # -> [batch_size * num_chunks, in_channels, lc+chunk_size+rpad]
+            out = out.flatten(start_dim=0, end_dim=1)
+
+            # TODO: experiment around reflect padding, which is difficult
+            # because small chunks have too little time steps to reflect from
+
+            # let's keep backwards compat by pointing at the weights from the
+            # already declared Conv1d.
+            #
+            # still reusing the above example, the convolution will be applied,
+            # with the padding truncated on both ends. the following example
+            # shows the letter corresponding to the input frame on which the
+            # convolution was centered.
+            #
+            # as you can see, the sum of lengths of all chunks is equal to our
+            # input sequence length + `final_right_padding`.
+            #
+            # chunk #1:  aaaa
+            # chunk #2:      bbbb
+            # chunk #3:          cc00
+
+            # -> [batch_size * num_chunks, out_channels, chunk_size]
+            out = F.conv1d(
+                out,
+                weight=self.conv.weight,
+                bias=self.conv.bias,
+                stride=self.conv.stride,
+                padding=0,
+                dilation=self.conv.dilation,
+                groups=self.conv.groups,
+            )
+
+            # -> [batch_size * num_chunks, chunk_size, out_channels]
+            out = out.transpose(1, 2)
+
+            out = self.after_conv(out)
+
+            # -> [batch_size, num_chunks, chunk_size, out_channels]
+            out = torch.unflatten(out, dim=0, sizes=(batch_size, -1))
+
+            # -> [batch_size, t + final_right_padding, out_channels]
+            out = torch.flatten(out, start_dim=1, end_dim=2)
+
+            # -> [batch_size, t, out_channels]
+            if final_right_padding > 0:
+                out = out[:, :-final_right_padding, :]
+        else:
+            out = self.layer_norm(x)
+            out = out.transpose(1, 2)
+            out = self.bottleneck(out)
+            out = self.conv(out)
+
+            if self.causal:
+                # chomp
+                out = out[..., : -self.padding]
+
+            out = out.transpose(1, 2)
+            out = self.after_conv(out)
+
+        if mask is not None:
+            out.masked_fill_(mask, 0.0)
+
+        return out
+
+
+class ConformerEncoderLayer(nn.Module):
+    """This is an implementation of Conformer encoder layer.
+
+    Arguments
+    ---------
+    d_model : int
+        The expected size of the input embedding.
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    nhead : int
+        Number of attention heads.
+    kernel_size : int, optional
+        Kernel size of convolution model.
+    kdim : int, optional
+        Dimension of the key.
+    vdim : int, optional
+        Dimension of the value.
+    activation: torch.nn.Module
+         Activation function used in each Conformer layer.
+    bias : bool, optional
+        Whether  convolution module.
+    dropout : int, optional
+        Dropout for the encoder.
+    causal : bool, optional
+        Whether the convolutions should be causal or not.
+    attention_type : str, optional
+        type of attention layer, e.g. regularMHA for regular MultiHeadAttention.
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> pos_embs = torch.rand((1, 2*60-1, 512))
+    >>> net = ConformerEncoderLayer(d_ffn=512, nhead=8, d_model=512, kernel_size=3)
+    >>> output = net(x, pos_embs=pos_embs)
+    >>> output[0].shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        d_model,
+        d_ffn,
+        nhead,
+        kernel_size=31,
+        kdim=None,
+        vdim=None,
+        activation=Swish,
+        bias=True,
+        dropout=0.0,
+        causal=False,
+        attention_type="RelPosMHAXL",
+    ):
+        super().__init__()
+
+        if attention_type == "regularMHA":
+            self.mha_layer = MultiheadAttention(
+                nhead=nhead,
+                d_model=d_model,
+                dropout=dropout,
+                kdim=kdim,
+                vdim=vdim,
+            )
+        elif attention_type == "RelPosMHAXL":
+            # transformerXL style positional encoding
+            self.mha_layer = RelPosMHAXL(
+                num_heads=nhead,
+                embed_dim=d_model,
+                dropout=dropout,
+                mask_pos_future=causal,
+            )
+        elif attention_type == "hypermixing":
+            self.mha_layer = HyperMixing(
+                input_output_dim=d_model,
+                hypernet_size=d_ffn,
+                tied=False,
+                num_heads=nhead,
+                fix_tm_hidden_size=False,
+            )
+
+        self.convolution_module = ConvolutionModule(
+            d_model, kernel_size, bias, activation, dropout, causal=causal
+        )
+
+        self.ffn_module1 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.ffn_module2 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: torch.Tensor = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """
+        Arguments
+        ----------
+        src : torch.Tensor
+            The sequence to the encoder layer.
+        src_mask : torch.Tensor, optional
+            The mask for the src sequence.
+        src_key_padding_mask : torch.Tensor, optional
+            The mask for the src keys per batch.
+        pos_embs: torch.Tensor, torch.nn.Module, optional
+            Module or tensor containing the input sequence positional embeddings
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming,
+            specifically involved here to apply Dynamic Chunk Convolution to
+            the convolution module.
+        """
+        conv_mask: Optional[torch.Tensor] = None
+        if src_key_padding_mask is not None:
+            conv_mask = src_key_padding_mask.unsqueeze(-1)
+        # ffn module
+        x = x + 0.5 * self.ffn_module1(x)
+        # multi-head attention module
+        skip = x
+        x = self.norm1(x)
+
+        x, self_attn = self.mha_layer(
+            x,
+            x,
+            x,
+            attn_mask=src_mask,
+            key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs,
+        )
+        x = x + skip
+        # convolution module
+        x = x + self.convolution_module(
+            x, conv_mask, dynchunktrain_config=dynchunktrain_config
+        )
+        # ffn module
+        x = self.norm2(x + 0.5 * self.ffn_module2(x))
+        return x, self_attn
+
+    def forward_streaming(
+        self,
+        x,
+        context: ConformerEncoderLayerStreamingContext,
+        pos_embs: torch.Tensor = None,
+    ):
+        """Conformer layer streaming forward (typically for
+        DynamicChunkTraining-trained models), which is to be used at inference
+        time. Relies on a mutable context object as initialized by
+        `make_streaming_context` that should be used across chunks.
+        Invoked by `ConformerEncoder.forward_streaming`.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input tensor for this layer. Batching is supported as long as you
+            keep the context consistent.
+        context : ConformerEncoderStreamingContext
+            Mutable streaming context; the same object should be passed across
+            calls.
+        pos_embs : torch.Tensor, optional
+            Positional embeddings, if used.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output tensor.
+        self_attn : list
+            List of self attention values.
+        """
+
+        orig_len = x.shape[-2]
+        # ffn module
+        x = x + 0.5 * self.ffn_module1(x)
+
+        # TODO: make the approach for MHA left context more efficient.
+        # currently, this saves the inputs to the MHA.
+        # the naive approach is suboptimal in a few ways, namely that the
+        # outputs for this left padding is being re-computed even though we
+        # discard them immediately after.
+
+        # left pad `x` with our MHA left context
+        if context.mha_left_context is not None:
+            x = torch.cat((context.mha_left_context, x), dim=1)
+
+        # compute new MHA left context for the next call to our function
+        if context.mha_left_context_size > 0:
+            context.mha_left_context = x[
+                ..., -context.mha_left_context_size :, :
+            ]
+
+        # multi-head attention module
+        skip = x
+        x = self.norm1(x)
+
+        x, self_attn = self.mha_layer(
+            x,
+            x,
+            x,
+            attn_mask=None,
+            key_padding_mask=None,
+            pos_embs=pos_embs,
+        )
+        x = x + skip
+
+        # truncate outputs corresponding to the MHA left context (we only care
+        # about our chunk's outputs); see above to-do
+        x = x[..., -orig_len:, :]
+
+        if context.dcconv_left_context is not None:
+            x = torch.cat((context.dcconv_left_context, x), dim=1)
+
+        # compute new DCConv left context for the next call to our function
+        context.dcconv_left_context = x[
+            ..., -self.convolution_module.padding :, :
+        ]
+
+        # convolution module
+        x = x + self.convolution_module(x)
+
+        # truncate outputs corresponding to the DCConv left context
+        x = x[..., -orig_len:, :]
+
+        # ffn module
+        x = self.norm2(x + 0.5 * self.ffn_module2(x))
+        return x, self_attn
+
+    def make_streaming_context(self, mha_left_context_size: int):
+        """Creates a blank streaming context for this encoding layer.
+
+        Arguments
+        ---------
+        mha_left_context_size : int
+            How many left frames should be saved and used as left context to the
+            current chunk when streaming
+
+        Returns
+        -------
+        ConformerEncoderLayerStreamingContext
+        """
+        return ConformerEncoderLayerStreamingContext(
+            mha_left_context_size=mha_left_context_size
+        )
+
+
+class ConformerEncoder(nn.Module):
+    """This class implements the Conformer encoder.
+
+    Arguments
+    ---------
+    num_layers : int
+        Number of layers.
+    d_model : int
+        Embedding dimension size.
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    nhead : int
+        Number of attention heads.
+    kernel_size : int, optional
+        Kernel size of convolution model.
+    kdim : int, optional
+        Dimension of the key.
+    vdim : int, optional
+        Dimension of the value.
+    activation: torch.nn.Module
+         Activation function used in each Confomer layer.
+    bias : bool, optional
+        Whether  convolution module.
+    dropout : int, optional
+        Dropout for the encoder.
+    causal: bool, optional
+        Whether the convolutions should be causal or not.
+    attention_type: str, optional
+        type of attention layer, e.g. regularMHA for regular MultiHeadAttention.
+
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> pos_emb = torch.rand((1, 2*60-1, 512))
+    >>> net = ConformerEncoder(1, 512, 512, 8)
+    >>> output, _ = net(x, pos_embs=pos_emb)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        d_model,
+        d_ffn,
+        nhead,
+        kernel_size=31,
+        kdim=None,
+        vdim=None,
+        activation=Swish,
+        bias=True,
+        dropout=0.0,
+        causal=False,
+        attention_type="RelPosMHAXL",
+    ):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(
+            [
+                ConformerEncoderLayer(
+                    d_ffn=d_ffn,
+                    nhead=nhead,
+                    d_model=d_model,
+                    kdim=kdim,
+                    vdim=vdim,
+                    dropout=dropout,
+                    activation=activation,
+                    kernel_size=kernel_size,
+                    bias=bias,
+                    causal=causal,
+                    attention_type=attention_type,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.norm = LayerNorm(d_model, eps=1e-6)
+        self.attention_type = attention_type
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: Optional[torch.Tensor] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """
+        Arguments
+        ----------
+        src : torch.Tensor
+            The sequence to the encoder layer.
+        src_mask : torch.Tensor, optional
+            The mask for the src sequence.
+        src_key_padding_mask : torch.Tensor, optional
+            The mask for the src keys per batch.
+        pos_embs: torch.Tensor, torch.nn.Module,
+            Module or tensor containing the input sequence positional embeddings
+            If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
+            where S is the sequence length, and E is the embedding dimension.
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming,
+            specifically involved here to apply Dynamic Chunk Convolution to the
+            convolution module.
+        """
+        if self.attention_type == "RelPosMHAXL":
+            if pos_embs is None:
+                raise ValueError(
+                    "The chosen attention type for the Conformer is RelPosMHAXL. For this attention type, the positional embeddings are mandatory"
+                )
+
+        output = src
+        attention_lst = []
+        for enc_layer in self.layers:
+            output, attention = enc_layer(
+                output,
+                src_mask=src_mask,
+                src_key_padding_mask=src_key_padding_mask,
+                pos_embs=pos_embs,
+                dynchunktrain_config=dynchunktrain_config,
+            )
+            attention_lst.append(attention)
+        output = self.norm(output)
+
+        return output, attention_lst
+
+    def forward_streaming(
+        self,
+        src: torch.Tensor,
+        context: ConformerEncoderStreamingContext,
+        pos_embs: Optional[torch.Tensor] = None,
+    ):
+        """Conformer streaming forward (typically for
+        DynamicChunkTraining-trained models), which is to be used at inference
+        time. Relies on a mutable context object as initialized by
+        `make_streaming_context` that should be used across chunks.
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            Input tensor. Batching is supported as long as you keep the context
+            consistent.
+        context : ConformerEncoderStreamingContext
+            Mutable streaming context; the same object should be passed across
+            calls.
+        pos_embs : torch.Tensor, optional
+            Positional embeddings, if used.
+
+        Returns
+        -------
+        output : torch.Tensor
+            The output of the streaming conformer.
+        attention_lst : list
+            The attention values.
+        """
+
+        if self.attention_type == "RelPosMHAXL":
+            if pos_embs is None:
+                raise ValueError(
+                    "The chosen attention type for the Conformer is RelPosMHAXL. For this attention type, the positional embeddings are mandatory"
+                )
+
+        output = src
+        attention_lst = []
+        for i, enc_layer in enumerate(self.layers):
+            output, attention = enc_layer.forward_streaming(
+                output, pos_embs=pos_embs, context=context.layers[i]
+            )
+            attention_lst.append(attention)
+        output = self.norm(output)
+
+        return output, attention_lst
+
+    def make_streaming_context(self, dynchunktrain_config: DynChunkTrainConfig):
+        """Creates a blank streaming context for the encoder.
+
+        Arguments
+        ---------
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming
+
+        Returns
+        -------
+        ConformerEncoderStreamingContext
+        """
+        return ConformerEncoderStreamingContext(
+            dynchunktrain_config=dynchunktrain_config,
+            layers=[
+                layer.make_streaming_context(
+                    mha_left_context_size=dynchunktrain_config.left_context_size_frames()
+                )
+                for layer in self.layers
+            ],
+        )
+
+
+class ConformerDecoderLayer(nn.Module):
+    """This is an implementation of Conformer encoder layer.
+
+    Arguments
+    ---------
+    d_model : int
+        The expected size of the input embedding.
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    nhead : int
+        Number of attention heads.
+    kernel_size : int, optional
+        Kernel size of convolution model.
+    kdim : int, optional
+        Dimension of the key.
+    vdim : int, optional
+        Dimension of the value.
+    activation : torch.nn.Module, optional
+         Activation function used in each Conformer layer.
+    bias : bool, optional
+        Whether  convolution module.
+    dropout : int, optional
+        Dropout for the encoder.
+    causal : bool, optional
+        Whether the convolutions should be causal or not.
+    attention_type : str, optional
+        type of attention layer, e.g. regularMHA for regular MultiHeadAttention.
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> pos_embs = torch.rand((1, 2*60-1, 512))
+    >>> net = ConformerEncoderLayer(d_ffn=512, nhead=8, d_model=512, kernel_size=3)
+    >>> output = net(x, pos_embs=pos_embs)
+    >>> output[0].shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        d_model,
+        d_ffn,
+        nhead,
+        kernel_size,
+        kdim=None,
+        vdim=None,
+        activation=Swish,
+        bias=True,
+        dropout=0.0,
+        causal=True,
+        attention_type="RelPosMHAXL",
+    ):
+        super().__init__()
+
+        if not causal:
+            warnings.warn(
+                "Decoder is not causal, in most applications it should be causal, you have been warned !"
+            )
+
+        if attention_type == "regularMHA":
+            self.mha_layer = MultiheadAttention(
+                nhead=nhead,
+                d_model=d_model,
+                dropout=dropout,
+                kdim=kdim,
+                vdim=vdim,
+            )
+        elif attention_type == "RelPosMHAXL":
+            # transformerXL style positional encoding
+            self.mha_layer = RelPosMHAXL(
+                num_heads=nhead,
+                embed_dim=d_model,
+                dropout=dropout,
+                mask_pos_future=causal,
+            )
+
+        self.convolution_module = ConvolutionModule(
+            d_model, kernel_size, bias, activation, dropout, causal=causal
+        )
+
+        self.ffn_module1 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.ffn_module2 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ---------
+        tgt: torch.Tensor
+            The sequence to the decoder layer.
+        memory: torch.Tensor
+            The sequence from the last layer of the encoder.
+        tgt_mask: torch.Tensor, optional, optional
+            The mask for the tgt sequence.
+        memory_mask: torch.Tensor, optional
+            The mask for the memory sequence.
+        tgt_key_padding_mask: torch.Tensor, optional
+            The mask for the tgt keys per batch.
+        memory_key_padding_mask: torch.Tensor, optional
+            The mask for the memory keys per batch.
+        pos_embs_tgt: torch.Tensor, torch.nn.Module, optional
+            Module or tensor containing the target sequence positional embeddings for each attention layer.
+        pos_embs_src: torch.Tensor, torch.nn.Module, optional
+            Module or tensor containing the source sequence positional embeddings for each attention layer.
+
+        Returns
+        -------
+        x: torch.Tensor
+            The output tensor
+        self_attn : torch.Tensor
+        self_attn : torch.Tensor
+            The self attention tensor
+        """
+        # ffn module
+        tgt = tgt + 0.5 * self.ffn_module1(tgt)
+        # multi-head attention module
+        skip = tgt
+        x = self.norm1(tgt)
+        x, self_attn = self.mha_layer(
+            x,
+            memory,
+            memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+            pos_embs=pos_embs_src,
+        )
+        x = x + skip
+        # convolution module
+        x = x + self.convolution_module(x)
+        # ffn module
+        x = self.norm2(x + 0.5 * self.ffn_module2(x))
+        return x, self_attn, self_attn
+
+
+class ConformerDecoder(nn.Module):
+    """This class implements the Transformer decoder.
+
+    Arguments
+    ---------
+    num_layers: int
+        Number of layers.
+    nhead: int
+        Number of attention heads.
+    d_ffn: int
+        Hidden size of self-attention Feed Forward layer.
+    d_model: int
+        Embedding dimension size.
+    kdim: int, optional
+        Dimension for key.
+    vdim: int, optional
+        Dimension for value.
+    dropout: float, optional
+        Dropout rate.
+    activation: torch.nn.Module, optional
+        Activation function used after non-bottleneck conv layer.
+    kernel_size : int, optional
+        Kernel size of convolutional layer.
+    bias : bool, optional
+        Whether  convolution module.
+    causal: bool, optional
+        Whether the convolutions should be causal or not.
+    attention_type: str, optional
+        type of attention layer, e.g. regularMHA for regular MultiHeadAttention.
+
+
+    Example
+    -------
+    >>> src = torch.rand((8, 60, 512))
+    >>> tgt = torch.rand((8, 60, 512))
+    >>> net = ConformerDecoder(1, 8, 1024, 512, attention_type="regularMHA")
+    >>> output, _, _ = net(tgt, src)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        nhead,
+        d_ffn,
+        d_model,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        activation=Swish,
+        kernel_size=3,
+        bias=True,
+        causal=True,
+        attention_type="RelPosMHAXL",
+    ):
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                ConformerDecoderLayer(
+                    d_ffn=d_ffn,
+                    nhead=nhead,
+                    d_model=d_model,
+                    kdim=kdim,
+                    vdim=vdim,
+                    dropout=dropout,
+                    activation=activation,
+                    kernel_size=kernel_size,
+                    bias=bias,
+                    causal=causal,
+                    attention_type=attention_type,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ---------
+        tgt: torch.Tensor
+            The sequence to the decoder layer.
+        memory: torch.Tensor
+            The sequence from the last layer of the encoder.
+        tgt_mask: torch.Tensor, optional, optional
+            The mask for the tgt sequence.
+        memory_mask: torch.Tensor, optional
+            The mask for the memory sequence.
+        tgt_key_padding_mask : torch.Tensor, optional
+            The mask for the tgt keys per batch.
+        memory_key_padding_mask : torch.Tensor, optional
+            The mask for the memory keys per batch.
+        pos_embs_tgt: torch.Tensor, torch.nn.Module, optional
+            Module or tensor containing the target sequence positional embeddings for each attention layer.
+        pos_embs_src: torch.Tensor, torch.nn.Module, optional
+            Module or tensor containing the source sequence positional embeddings for each attention layer.
+
+        Returns
+        -------
+        output: torch.Tensor
+            Conformer decoder output.
+        self_attns : list
+            Location of self attentions.
+        multihead_attns : list
+            Location of multihead attentions.
+        """
+        output = tgt
+        self_attns, multihead_attns = [], []
+        for dec_layer in self.layers:
+            output, self_attn, multihead_attn = dec_layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                pos_embs_tgt=pos_embs_tgt,
+                pos_embs_src=pos_embs_src,
+            )
+            self_attns.append(self_attn)
+            multihead_attns.append(multihead_attn)
+        output = self.norm(output)
+
+        return output, self_attns, multihead_attns

modules/Conmamba.py

@@ -0,0 +1,607 @@
+"""ConMamba encoder and Mamba decoder implementation.
+
+Authors
+-------
+* Xilin Jiang 2024
+"""
+
+import warnings
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import speechbrain as sb
+from speechbrain.nnet.activations import Swish
+from speechbrain.nnet.attention import (
+    MultiheadAttention,
+    PositionalwiseFeedForward,
+    RelPosMHAXL,
+)
+from speechbrain.nnet.hypermixing import HyperMixing
+from speechbrain.nnet.normalization import LayerNorm
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+
+# Mamba
+from mamba_ssm import Mamba
+from modules.mamba.bimamba import Mamba as BiMamba 
+
+
+class ConvolutionModule(nn.Module):
+    """This is an implementation of convolution module in Conmamba.
+    """
+
+    def __init__(
+        self,
+        input_size,
+        kernel_size=31,
+        bias=True,
+        activation=Swish,
+        dropout=0.0,
+        causal=False,
+        dilation=1,
+    ):
+        super().__init__()
+
+        self.kernel_size = kernel_size
+        self.causal = causal
+        self.dilation = dilation
+
+        if self.causal:
+            self.padding = (kernel_size - 1) * 2 ** (dilation - 1)
+        else:
+            self.padding = (kernel_size - 1) * 2 ** (dilation - 1) // 2
+
+        self.layer_norm = nn.LayerNorm(input_size)
+        self.bottleneck = nn.Sequential(
+            # pointwise
+            nn.Conv1d(
+                input_size, 2 * input_size, kernel_size=1, stride=1, bias=bias
+            ),
+            nn.GLU(dim=1),
+        )
+        # depthwise
+        self.conv = nn.Conv1d(
+            input_size,
+            input_size,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=self.padding,
+            dilation=dilation,
+            groups=input_size,
+            bias=bias,
+        )
+
+        # BatchNorm in the original Conformer replaced with a LayerNorm due to
+        # https://github.com/speechbrain/speechbrain/pull/1329
+        # see discussion
+        # https://github.com/speechbrain/speechbrain/pull/933#issuecomment-1033367884
+
+        self.after_conv = nn.Sequential(
+            nn.LayerNorm(input_size),
+            activation(),
+            # pointwise
+            nn.Linear(input_size, input_size, bias=bias),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """Applies the convolution to an input tensor `x`.
+        """
+
+        if dynchunktrain_config is not None:
+            # chances are chunking+causal is unintended; i don't know where it
+            # may make sense, but if it does to you, feel free to implement it.
+            assert (
+                not self.causal
+            ), "Chunked convolution not supported with causal padding"
+
+            assert (
+                self.dilation == 1
+            ), "Current DynChunkTrain logic does not support dilation != 1"
+
+            # in a causal convolution, which is not the case here, an output
+            # frame would never be able to depend on a input frame from any
+            # point in the future.
+
+            # but with the dynamic chunk convolution, we instead use a "normal"
+            # convolution but where, for any output frame, the future beyond the
+            # "current" chunk gets masked.
+            # see the paper linked in the documentation for details.
+
+            chunk_size = dynchunktrain_config.chunk_size
+            batch_size = x.shape[0]
+
+            # determine the amount of padding we need to insert at the right of
+            # the last chunk so that all chunks end up with the same size.
+            if x.shape[1] % chunk_size != 0:
+                final_right_padding = chunk_size - (x.shape[1] % chunk_size)
+            else:
+                final_right_padding = 0
+
+            # -> [batch_size, t, in_channels]
+            out = self.layer_norm(x)
+
+            # -> [batch_size, in_channels, t] for the CNN
+            out = out.transpose(1, 2)
+
+            # -> [batch_size, in_channels, t] (pointwise)
+            out = self.bottleneck(out)
+
+            # -> [batch_size, in_channels, lc+t+final_right_padding]
+            out = F.pad(out, (self.padding, final_right_padding), value=0)
+
+            # now, make chunks with left context.
+            # as a recap to what the above padding and this unfold do, consider
+            # each a/b/c letter represents a frame as part of chunks a, b, c.
+            # consider a chunk size of 4 and a kernel size of 5 (padding=2):
+            #
+            # input seq: 00aaaabbbbcc00
+            # chunk #1:  00aaaa
+            # chunk #2:      aabbbb
+            # chunk #3:          bbcc00
+            #
+            # a few remarks here:
+            # - the left padding gets inserted early so that the unfold logic
+            #   works trivially
+            # - the right 0-padding got inserted as the number of time steps
+            #   could not be evenly split in `chunk_size` chunks
+
+            # -> [batch_size, in_channels, num_chunks, lc+chunk_size]
+            out = out.unfold(2, size=chunk_size + self.padding, step=chunk_size)
+
+            # as we manually disable padding in the convolution below, we insert
+            # right 0-padding to the chunks, e.g. reusing the above example:
+            #
+            # chunk #1:  00aaaa00
+            # chunk #2:      aabbbb00
+            # chunk #3:          bbcc0000
+
+            # -> [batch_size, in_channels, num_chunks, lc+chunk_size+rpad]
+            out = F.pad(out, (0, self.padding), value=0)
+
+            # the transpose+flatten effectively flattens chunks into the batch
+            # dimension to be processed into the time-wise convolution. the
+            # chunks will later on be unflattened.
+
+            # -> [batch_size, num_chunks, in_channels, lc+chunk_size+rpad]
+            out = out.transpose(1, 2)
+
+            # -> [batch_size * num_chunks, in_channels, lc+chunk_size+rpad]
+            out = out.flatten(start_dim=0, end_dim=1)
+
+            # TODO: experiment around reflect padding, which is difficult
+            # because small chunks have too little time steps to reflect from
+
+            # let's keep backwards compat by pointing at the weights from the
+            # already declared Conv1d.
+            #
+            # still reusing the above example, the convolution will be applied,
+            # with the padding truncated on both ends. the following example
+            # shows the letter corresponding to the input frame on which the
+            # convolution was centered.
+            #
+            # as you can see, the sum of lengths of all chunks is equal to our
+            # input sequence length + `final_right_padding`.
+            #
+            # chunk #1:  aaaa
+            # chunk #2:      bbbb
+            # chunk #3:          cc00
+
+            # -> [batch_size * num_chunks, out_channels, chunk_size]
+            out = F.conv1d(
+                out,
+                weight=self.conv.weight,
+                bias=self.conv.bias,
+                stride=self.conv.stride,
+                padding=0,
+                dilation=self.conv.dilation,
+                groups=self.conv.groups,
+            )
+
+            # -> [batch_size * num_chunks, chunk_size, out_channels]
+            out = out.transpose(1, 2)
+
+            out = self.after_conv(out)
+
+            # -> [batch_size, num_chunks, chunk_size, out_channels]
+            out = torch.unflatten(out, dim=0, sizes=(batch_size, -1))
+
+            # -> [batch_size, t + final_right_padding, out_channels]
+            out = torch.flatten(out, start_dim=1, end_dim=2)
+
+            # -> [batch_size, t, out_channels]
+            if final_right_padding > 0:
+                out = out[:, :-final_right_padding, :]
+        else:
+            out = self.layer_norm(x)
+            out = out.transpose(1, 2)
+            out = self.bottleneck(out)
+            out = self.conv(out)
+
+            if self.causal:
+                # chomp
+                out = out[..., : -self.padding]
+
+            out = out.transpose(1, 2)
+            out = self.after_conv(out)
+
+        if mask is not None:
+            out.masked_fill_(mask, 0.0)
+
+        return out
+
+
+class ConmambaEncoderLayer(nn.Module):
+    """This is an implementation of Conmamba encoder layer.
+    """
+
+    def __init__(
+        self,
+        d_model,
+        d_ffn,
+        kernel_size=31,
+        activation=Swish,
+        bias=True,
+        dropout=0.0,
+        causal=False,
+        mamba_config=None
+    ):
+        super().__init__()
+        assert mamba_config != None
+
+        bidirectional = mamba_config.pop('bidirectional')
+        if causal or (not bidirectional):
+            self.mamba = Mamba(
+                d_model=d_model,
+                **mamba_config
+            )
+        else:
+            self.mamba = BiMamba(
+                d_model=d_model,
+                bimamba_type='v2',
+                **mamba_config
+            )
+        mamba_config['bidirectional'] = bidirectional
+
+        self.convolution_module = ConvolutionModule(
+            d_model, kernel_size, bias, activation, dropout, causal=causal
+        )
+
+        self.ffn_module1 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.ffn_module2 = nn.Sequential(
+            nn.LayerNorm(d_model),
+            PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            ),
+            nn.Dropout(dropout),
+        )
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: torch.Tensor = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        conv_mask: Optional[torch.Tensor] = None
+        if src_key_padding_mask is not None:
+            conv_mask = src_key_padding_mask.unsqueeze(-1)
+
+        conv_mask = None
+
+        # ffn module
+        x = x + 0.5 * self.ffn_module1(x)
+        # mamba module
+        skip = x
+        x = self.norm1(x)
+        x = self.mamba(x)
+        x = x + skip
+        # convolution module
+        x = x + self.convolution_module(
+            x, conv_mask, dynchunktrain_config=dynchunktrain_config
+        )
+        # ffn module
+        x = self.norm2(x + 0.5 * self.ffn_module2(x))
+        return x
+
+
+class ConmambaEncoder(nn.Module):
+    """This class implements the Conmamba encoder.
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        d_model,
+        d_ffn,
+        kernel_size=31,
+        activation=Swish,
+        bias=True,
+        dropout=0.0,
+        causal=False,
+        mamba_config=None
+    ):
+        super().__init__()
+        print(f'dropout={str(dropout)} is not used in Mamba.')
+
+        self.layers = torch.nn.ModuleList(
+            [
+                ConmambaEncoderLayer(
+                    d_model=d_model,
+                    d_ffn=d_ffn,
+                    dropout=dropout,
+                    activation=activation,
+                    kernel_size=kernel_size,
+                    bias=bias,
+                    causal=causal,
+                    mamba_config=mamba_config,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.norm = LayerNorm(d_model, eps=1e-6)
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: Optional[torch.Tensor] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """
+        Arguments
+        ----------
+        src : torch.Tensor
+            The sequence to the encoder layer.
+        src_mask : torch.Tensor, optional
+            The mask for the src sequence.
+        src_key_padding_mask : torch.Tensor, optional
+            The mask for the src keys per batch.
+        pos_embs: torch.Tensor, torch.nn.Module,
+            Module or tensor containing the input sequence positional embeddings
+            If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
+            where S is the sequence length, and E is the embedding dimension.
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming,
+            specifically involved here to apply Dynamic Chunk Convolution to the
+            convolution module.
+        """
+
+        output = src
+        for enc_layer in self.layers:
+            output = enc_layer(
+                output,
+                src_mask=src_mask,
+                src_key_padding_mask=src_key_padding_mask,
+                pos_embs=pos_embs,
+                dynchunktrain_config=dynchunktrain_config,
+            )
+        output = self.norm(output)
+
+        return output, None
+
+
+class MambaDecoderLayer(nn.Module):
+    """This class implements the Mamba decoder layer.
+    """
+
+    def __init__(
+        self,
+        d_model,
+        d_ffn,
+        activation=nn.ReLU,
+        dropout=0.0,
+        normalize_before=False,
+        mamba_config=None
+    ):
+        super().__init__()
+
+        assert mamba_config != None
+
+        bidirectional = mamba_config.pop('bidirectional')
+
+        self.self_mamba = Mamba(
+            d_model=d_model,
+            **mamba_config
+        )
+
+        self.cross_mamba = Mamba(
+            d_model=d_model,
+            **mamba_config
+        )
+
+        mamba_config['bidirectional'] = bidirectional
+
+        self.pos_ffn = sb.nnet.attention.PositionalwiseFeedForward(
+            d_ffn=d_ffn,
+            input_size=d_model,
+            dropout=dropout,
+            activation=activation,
+        )
+
+        # normalization layers
+        self.norm1 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.norm2 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.norm3 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+        self.dropout3 = torch.nn.Dropout(dropout)
+
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ----------
+        tgt: torch.Tensor
+            The sequence to the decoder layer (required).
+        memory: torch.Tensor
+            The sequence from the last layer of the encoder (required).
+        tgt_mask: torch.Tensor
+            The mask for the tgt sequence (optional).
+        memory_mask: torch.Tensor
+            The mask for the memory sequence (optional).
+        tgt_key_padding_mask: torch.Tensor
+            The mask for the tgt keys per batch (optional).
+        memory_key_padding_mask: torch.Tensor
+            The mask for the memory keys per batch (optional).
+        pos_embs_tgt: torch.Tensor
+            The positional embeddings for the target (optional).
+        pos_embs_src: torch.Tensor
+            The positional embeddings for the source (optional).
+        """
+        if self.normalize_before:
+            tgt1 = self.norm1(tgt)
+        else:
+            tgt1 = tgt
+
+        # Mamba over the target sequence
+        tgt2 = self.self_mamba(tgt1)
+
+        # add & norm
+        tgt = tgt + self.dropout1(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm1(tgt)
+
+        if self.normalize_before:
+            tgt1 = self.norm2(tgt)
+        else:
+            tgt1 = tgt
+
+        # Mamba over key=value + query
+        # and only take the last len(query) tokens
+        tgt2 = self.cross_mamba(torch.cat([memory, tgt1], dim=1))[:, -tgt1.shape[1]:]
+        
+        # add & norm
+        tgt = tgt + self.dropout2(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm2(tgt)
+
+        if self.normalize_before:
+            tgt1 = self.norm3(tgt)
+        else:
+            tgt1 = tgt
+
+        tgt2 = self.pos_ffn(tgt1)
+
+        # add & norm
+        tgt = tgt + self.dropout3(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm3(tgt)
+
+        return tgt, None, None
+
+
+class MambaDecoder(nn.Module):
+    """This class implements the Mamba decoder.
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        d_model,
+        d_ffn,
+        activation=nn.ReLU,
+        dropout=0.0,
+        normalize_before=False,
+        mamba_config=None
+    ):
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                MambaDecoderLayer(
+                    d_model=d_model,
+                    d_ffn=d_ffn,
+                    activation=activation,
+                    dropout=dropout,
+                    normalize_before=normalize_before,
+                    mamba_config=mamba_config
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ----------
+        tgt : torch.Tensor
+            The sequence to the decoder layer (required).
+        memory : torch.Tensor
+            The sequence from the last layer of the encoder (required).
+        tgt_mask : torch.Tensor
+            The mask for the tgt sequence (optional).
+        memory_mask : torch.Tensor
+            The mask for the memory sequence (optional).
+        tgt_key_padding_mask : torch.Tensor
+            The mask for the tgt keys per batch (optional).
+        memory_key_padding_mask : torch.Tensor
+            The mask for the memory keys per batch (optional).
+        pos_embs_tgt : torch.Tensor
+            The positional embeddings for the target (optional).
+        pos_embs_src : torch.Tensor
+            The positional embeddings for the source (optional).
+        """
+        output = tgt
+        for dec_layer in self.layers:
+            output, _, _ = dec_layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                pos_embs_tgt=pos_embs_tgt,
+                pos_embs_src=pos_embs_src,
+            )
+        output = self.norm(output)
+
+        return output, [None], [None]

modules/MambaASR.py

@@ -0,0 +1,588 @@
+"""Added ConMamba and Mamba
+
+Authors
+* Xilin Jiang 2024
+"""
+
+"""Transformer for ASR in the SpeechBrain style.
+
+Authors
+* Jianyuan Zhong 2020
+* Titouan Parcollet 2024
+* Luca Della Libera 2024
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch  # noqa 42
+from torch import nn
+
+from speechbrain.dataio.dataio import length_to_mask
+from modules.Transformer import (
+    NormalizedEmbedding,
+    TransformerInterface,
+    get_key_padding_mask,
+    get_lookahead_mask,
+)
+from speechbrain.nnet.activations import Swish
+from speechbrain.nnet.containers import ModuleList
+from speechbrain.nnet.linear import Linear
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+
+
+@dataclass
+class TransformerASRStreamingContext:
+    """Streaming metadata and state for a `TransformerASR` instance."""
+
+    dynchunktrain_config: DynChunkTrainConfig
+    """Dynamic Chunk Training configuration holding chunk size and context size
+    information."""
+
+    encoder_context: Any
+    """Opaque encoder context information. It is constructed by the encoder's
+    `make_streaming_context` method and is passed to the encoder when using
+    `encode_streaming`.
+    """
+
+
+def make_transformer_src_mask(
+    src: torch.Tensor,
+    causal: bool = False,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+) -> Optional[torch.Tensor]:
+    """Prepare the source transformer mask that restricts which frames can
+    attend to which frames depending on causal or other simple restricted
+    attention methods.
+
+    Arguments
+    ---------
+    src: torch.Tensor
+        The source tensor to build a mask from. The contents of the tensor are
+        not actually used currently; only its shape and other metadata (e.g.
+        device).
+    causal: bool
+        Whether strict causality shall be used. Frames will not be able to
+        attend to any future frame.
+    dynchunktrain_config: DynChunkTrainConfig, optional
+        Dynamic Chunk Training configuration. This implements a simple form of
+        chunkwise attention. Incompatible with `causal`.
+
+    Returns
+    -------
+    torch.Tensor
+        A boolean mask Tensor of shape (timesteps, timesteps).
+    """
+    if causal:
+        assert dynchunktrain_config is None
+        return get_lookahead_mask(src)
+
+    if dynchunktrain_config is None:
+        return
+
+    # The following is not really the sole source used to implement this,
+    # but it helps introduce the concept.
+    # ref: Unified Streaming and Non-streaming Two-pass End-to-end Model for Speech Recognition
+    # https://arxiv.org/pdf/2012.05481.pdf
+    timesteps = src.size(1)
+
+    # Mask the future at the right of each chunk
+    chunk_size = dynchunktrain_config.chunk_size
+    num_chunks = timesteps // chunk_size
+    timestep_idx = torch.arange(timesteps, device=src.device)
+    mask_idx = torch.arange(
+        chunk_size, chunk_size * (num_chunks + 2), chunk_size, device=src.device
+    ).repeat_interleave(chunk_size)[:timesteps]
+    src_mask = timestep_idx[None] >= mask_idx[:, None]
+
+    # Mask the past at the left of each chunk (accounting for left context)
+    # only relevant if using left context
+    if not dynchunktrain_config.is_infinite_left_context():
+        num_left_chunks = dynchunktrain_config.left_context_size
+        mask_idx -= chunk_size * (num_left_chunks + 1)
+        src_mask += timestep_idx[None] < mask_idx[:, None]
+
+    return src_mask
+
+
+def make_transformer_src_tgt_masks(
+    src,
+    tgt=None,
+    wav_len=None,
+    pad_idx=0,
+    causal: bool = False,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+):
+    """This function generates masks for training the transformer model,
+    opinionated for an ASR context with encoding masks and, optionally, decoding
+    masks (if specifying `tgt`).
+
+    Arguments
+    ---------
+    src : torch.Tensor
+        The sequence to the encoder (required).
+    tgt : torch.Tensor
+        The sequence to the decoder.
+    wav_len : torch.Tensor
+        The lengths of the inputs.
+    pad_idx : int
+        The index for <pad> token (default=0).
+    causal: bool
+        Whether strict causality shall be used. See `make_asr_src_mask`
+    dynchunktrain_config: DynChunkTrainConfig, optional
+        Dynamic Chunk Training configuration. See `make_asr_src_mask`
+
+    Returns
+    -------
+    src_key_padding_mask : torch.Tensor
+        Key padding mask for ignoring padding
+    tgt_key_padding_mask : torch.Tensor
+        Key padding mask for ignoring padding
+    src_mask : torch.Tensor
+        Mask for ignoring invalid (e.g. future) timesteps
+    tgt_mask : torch.Tensor
+        Mask for ignoring invalid (e.g. future) timesteps
+    """
+    src_key_padding_mask = None
+
+    # mask out audio beyond the length of audio for each batch
+    if wav_len is not None:
+        abs_len = torch.round(wav_len * src.shape[1])
+        src_key_padding_mask = ~length_to_mask(abs_len).bool()
+
+    # mask out the source
+    src_mask = make_transformer_src_mask(
+        src, causal=causal, dynchunktrain_config=dynchunktrain_config
+    )
+
+    # If no decoder in the transformer...
+    if tgt is not None:
+        tgt_key_padding_mask = get_key_padding_mask(tgt, pad_idx=pad_idx)
+        tgt_mask = get_lookahead_mask(tgt)
+    else:
+        tgt_key_padding_mask = None
+        tgt_mask = None
+
+    return src_key_padding_mask, tgt_key_padding_mask, src_mask, tgt_mask
+
+
+class MambaASR(TransformerInterface):
+    def __init__(
+        self,
+        tgt_vocab,
+        input_size,
+        d_model=512,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        d_ffn=2048,
+        activation=nn.ReLU,
+        normalize_before=False,
+        kernel_size: Optional[int] = 31,
+        bias: Optional[bool] = True,
+        encoder_module: Optional[str] = "conmamba",
+        decoder_module: Optional[str] = "mamba",
+        max_length: Optional[int] = 2500,
+        causal: Optional[bool] = False,
+        csgu_linear_units: Optional[int] = 3072,
+        gate_activation: Optional[nn.Module] = nn.Identity,
+        use_linear_after_conv: Optional[bool] = False,
+        mamba_config=None
+    ):
+        super().__init__(
+            d_model=d_model,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            d_ffn=d_ffn,
+            activation=activation,
+            normalize_before=normalize_before,
+            kernel_size=kernel_size,
+            bias=bias,
+            encoder_module=encoder_module,
+            decoder_module=decoder_module,
+            max_length=max_length,
+            causal=causal,
+            csgu_linear_units=csgu_linear_units,
+            gate_activation=gate_activation,
+            use_linear_after_conv=use_linear_after_conv,
+            mamba_config=mamba_config
+        )
+
+        self.custom_src_module = ModuleList(
+            Linear(
+                input_size=input_size,
+                n_neurons=d_model,
+                bias=True,
+                combine_dims=False,
+            ),
+        )
+
+        self.num_decoder_layers = num_decoder_layers
+        if num_decoder_layers > 0:
+            self.custom_tgt_module = ModuleList(
+                NormalizedEmbedding(d_model, tgt_vocab)
+            )
+
+        # reset parameters using xavier_normal_
+        self._init_params()
+
+    def forward(self, src, tgt, wav_len=None, pad_idx=0):
+        """
+        Arguments
+        ----------
+        src : torch.Tensor
+            The sequence to the encoder.
+        tgt : torch.Tensor
+            The sequence to the decoder.
+        wav_len: torch.Tensor, optional
+            Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
+        pad_idx : int, optional
+            The index for <pad> token (default=0).
+        """
+
+        # reshape the src vector to [Batch, Time, Fea] is a 4d vector is given
+        if src.ndim == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        (
+            src_key_padding_mask,
+            tgt_key_padding_mask,
+            src_mask,
+            tgt_mask,
+        ) = make_transformer_src_tgt_masks(
+            src, tgt, wav_len, causal=self.causal, pad_idx=pad_idx
+        )
+        
+        src = self.custom_src_module(src)
+        # add pos encoding to queries if are sinusoidal ones else
+        if self.attention_type == "hypermixing":
+            pos_embs_encoder = None
+        elif self.attention_type == "RelPosMHAXL":
+            pos_embs_encoder = self.positional_encoding(src)
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(src)  # add the encodings here
+            pos_embs_encoder = None
+
+        encoder_out, _ = self.encoder(
+            src=src,
+            src_mask=src_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs_encoder,
+        )
+
+        if self.num_decoder_layers > 0:
+            tgt = self.custom_tgt_module(tgt)
+
+            if self.attention_type == "RelPosMHAXL":
+                tgt = tgt + self.positional_encoding_decoder(tgt)
+                pos_embs_encoder = None  # self.positional_encoding(src)
+                pos_embs_target = None
+            elif (
+                self.positional_encoding_type == "fixed_abs_sine"
+                or self.attention_type == "hypermixing"
+            ):
+                tgt = tgt + self.positional_encoding(tgt)
+                pos_embs_target = None
+                pos_embs_encoder = None
+
+            decoder_out, _, _ = self.decoder(
+                tgt=tgt,
+                memory=encoder_out,
+                memory_mask=None,
+                tgt_mask=tgt_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=src_key_padding_mask,
+                pos_embs_tgt=pos_embs_target,
+                pos_embs_src=pos_embs_encoder,
+            )
+
+        else:
+            decoder_out = None
+
+        return encoder_out, decoder_out
+
+    @torch.no_grad()
+    def decode(self, tgt, encoder_out, enc_len=None):
+        """This method implements a decoding step for the transformer model.
+
+        Arguments
+        ---------
+        tgt : torch.Tensor
+            The sequence to the decoder.
+        encoder_out : torch.Tensor
+            Hidden output of the encoder.
+        enc_len : torch.LongTensor
+            The actual length of encoder states.
+
+        Returns
+        -------
+        prediction
+        """
+        tgt_mask = get_lookahead_mask(tgt)
+        src_key_padding_mask = None
+        if enc_len is not None:
+            src_key_padding_mask = (1 - length_to_mask(enc_len)).bool()
+
+        if self.num_decoder_layers > 0:
+            tgt = self.custom_tgt_module(tgt)
+        if self.attention_type == "RelPosMHAXL":
+            tgt = tgt + self.positional_encoding_decoder(tgt)
+            pos_embs_encoder = None  # self.positional_encoding(src)
+            pos_embs_target = None
+        elif (
+            self.positional_encoding_type == "fixed_abs_sine"
+            or self.attention_type == "hypermixing"
+        ):
+            tgt = tgt + self.positional_encoding(tgt)  # add the encodings here
+            pos_embs_target = None
+            pos_embs_encoder = None
+
+   
+        prediction, self_attns, multihead_attns = self.decoder(
+            tgt,
+            encoder_out,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_key_padding_mask,
+            pos_embs_tgt=pos_embs_target,
+            pos_embs_src=pos_embs_encoder,
+        )
+        return prediction, multihead_attns[-1]
+
+    def encode(
+        self,
+        src,
+        wav_len=None,
+        pad_idx=0,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """
+        Encoder forward pass
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence to the encoder.
+        wav_len : torch.Tensor, optional
+            Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
+        pad_idx : int
+            The index used for padding.
+        dynchunktrain_config : DynChunkTrainConfig
+            Dynamic chunking config.
+
+        Returns
+        -------
+        encoder_out : torch.Tensor
+        """
+        # reshape the src vector to [Batch, Time, Fea] if a 4d vector is given
+        if src.dim() == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        (
+            src_key_padding_mask,
+            _,
+            src_mask,
+            _,
+        ) = make_transformer_src_tgt_masks(
+            src,
+            None,
+            wav_len,
+            pad_idx=pad_idx,
+            causal=self.causal,
+            dynchunktrain_config=dynchunktrain_config,
+        )
+
+        src = self.custom_src_module(src)
+        if self.attention_type == "hypermixing":
+            pos_embs_source = None
+        elif self.attention_type == "RelPosMHAXL":
+            pos_embs_source = self.positional_encoding(src)
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(src)
+            pos_embs_source = None
+
+        encoder_out, _ = self.encoder(
+            src=src,
+            src_mask=src_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs_source,
+            dynchunktrain_config=dynchunktrain_config,
+        )
+
+        return encoder_out
+
+    def encode_streaming(self, src, context: TransformerASRStreamingContext):
+        """
+        Streaming encoder forward pass
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence (chunk) to the encoder.
+        context : TransformerASRStreamingContext
+            Mutable reference to the streaming context. This holds the state
+            needed to persist across chunk inferences and can be built using
+            `make_streaming_context`. This will get mutated by this function.
+
+        Returns
+        -------
+        Encoder output for this chunk.
+
+        Example
+        -------
+        >>> import torch
+        >>> from speechbrain.lobes.models.transformer.TransformerASR import TransformerASR
+        >>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+        >>> net = TransformerASR(
+        ...     tgt_vocab=100,
+        ...     input_size=64,
+        ...     d_model=64,
+        ...     nhead=8,
+        ...     num_encoder_layers=1,
+        ...     num_decoder_layers=0,
+        ...     d_ffn=128,
+        ...     attention_type="RelPosMHAXL",
+        ...     positional_encoding=None,
+        ...     encoder_module="conformer",
+        ...     normalize_before=True,
+        ...     causal=False,
+        ... )
+        >>> ctx = net.make_streaming_context(DynChunkTrainConfig(16, 1))
+        >>> src1 = torch.rand([8, 16, 64])
+        >>> src2 = torch.rand([8, 16, 64])
+        >>> out1 = net.encode_streaming(src1, ctx)
+        >>> out1.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 16, 64])
+        >>> out2 = net.encode_streaming(src2, ctx)
+        >>> out2.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 16, 64])
+        >>> combined_out = torch.concat((out1, out2), dim=1)
+        >>> combined_out.shape
+        torch.Size([8, 32, 64])
+        """
+
+        if src.dim() == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        # HACK: our problem here is that the positional_encoding is computed
+        # against the size of our source tensor, but we only know how many left
+        # context frames we're injecting to the encoder within the encoder
+        # context.
+        # so this workaround does just that.
+        #
+        # i'm not sure how this would be best refactored, but an option would be
+        # to let the encoder get the pos embedding itself and have a way to
+        # cache it.
+        #
+        # additionally, positional encoding functions take in a whole source
+        # tensor just to get its attributes (size, device, type) but this is
+        # sort of silly for the embeddings that don't need one.
+        # so we craft a dummy empty (uninitialized) tensor to help...
+        known_left_context = context.encoder_context.layers[0].mha_left_context
+        if known_left_context is None:
+            pos_encoding_dummy = src
+        else:
+            target_shape = list(src.shape)
+            target_shape[-2] += known_left_context.shape[-2]
+            pos_encoding_dummy = torch.empty(size=target_shape).to(src)
+
+        src = self.custom_src_module(src)
+        if self.attention_type == "RelPosMHAXL":
+            pos_embs_source = self.positional_encoding(pos_encoding_dummy)
+
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(pos_encoding_dummy)
+            pos_embs_source = None
+
+        encoder_out, _ = self.encoder.forward_streaming(
+            src=src, pos_embs=pos_embs_source, context=context.encoder_context
+        )
+        return encoder_out
+
+    def make_streaming_context(
+        self, dynchunktrain_config: DynChunkTrainConfig, encoder_kwargs={}
+    ):
+        """Creates a blank streaming context for this transformer and its
+        encoder.
+
+        Arguments
+        ---------
+        dynchunktrain_config : DynChunkTrainConfig
+            Runtime chunkwise attention configuration.
+        encoder_kwargs : dict
+            Parameters to be forward to the encoder's `make_streaming_context`.
+            Metadata required for the encoder could differ depending on the
+            encoder.
+
+        Returns
+        -------
+        TransformerASRStreamingContext
+        """
+        return TransformerASRStreamingContext(
+            dynchunktrain_config=dynchunktrain_config,
+            encoder_context=self.encoder.make_streaming_context(
+                dynchunktrain_config,
+                **encoder_kwargs,
+            ),
+        )
+
+    def _init_params(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                torch.nn.init.xavier_normal_(p)
+
+
+class EncoderWrapper(nn.Module):
+    """This is a wrapper of any ASR transformer encoder. By default, the
+    TransformerASR .forward() function encodes and decodes. With this wrapper
+    the .forward() function becomes .encode() only.
+
+    Important: The TransformerASR class must contain a .encode() function.
+
+    Arguments
+    ---------
+    transformer : sb.lobes.models.TransformerInterface
+        A Transformer instance that contains a .encode() function.
+    *args : tuple
+    **kwargs : dict
+        Arguments to forward to parent class.
+
+    Example
+    -------
+    >>> src = torch.rand([8, 120, 512])
+    >>> tgt = torch.randint(0, 720, [8, 120])
+    >>> net = TransformerASR(
+    ...     720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU
+    ... )
+    >>> encoder = EncoderWrapper(net)
+    >>> enc_out = encoder(src)
+    >>> enc_out.shape
+    torch.Size([8, 120, 512])
+    """
+
+    def __init__(self, transformer, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.transformer = transformer
+        self.make_streaming_context = self.transformer.make_streaming_context
+
+    def forward(self, x, wav_lens=None, pad_idx=0, **kwargs):
+        """Processes the input tensor x and returns an output tensor."""
+        x = self.transformer.encode(x, wav_lens, pad_idx, **kwargs)
+        return x
+
+    def forward_streaming(self, x, context):
+        """Processes the input audio chunk tensor `x`, using and updating the
+        mutable encoder `context`"""
+        x = self.transformer.encode_streaming(x, context)
+        return x
+
+    def make_streaming_context(self, *args, **kwargs):
+        """Initializes a streaming context. Forwards all arguments to the
+        underlying transformer. See :meth:`speechbrain.lobes.models.transformer.TransformerASR.make_streaming_context`.
+        """
+        return self.transformer.make_streaming_context(*args, **kwargs)

modules/Transformer.py

@@ -0,0 +1,1085 @@
+"""Added ConMamba and Mamba
+
+Authors
+* Xilin Jiang 2024
+"""
+
+"""Transformer implementation in the SpeechBrain style.
+
+Authors
+* Jianyuan Zhong 2020
+* Samuele Cornell 2021
+"""
+
+import math
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import speechbrain as sb
+from speechbrain.nnet.activations import Swish
+from speechbrain.nnet.attention import RelPosEncXL
+from speechbrain.nnet.CNN import Conv1d
+
+from modules.Conformer import ConformerEncoder
+from modules.Conmamba import ConmambaEncoder, MambaDecoder
+
+
+class TransformerInterface(nn.Module):
+    """This is an interface for transformer model.
+    Users can modify the attributes and define the forward function as
+    needed according to their own tasks.
+    The architecture is based on the paper "Attention Is All You Need":
+    https://arxiv.org/pdf/1706.03762.pdf
+
+    Arguments
+    ---------
+    d_model: int
+        The number of expected features in the encoder/decoder inputs (default=512).
+    nhead: int
+        The number of heads in the multi-head attention models (default=8).
+    num_encoder_layers: int, optional
+        The number of encoder layers in1ì the encoder.
+    num_decoder_layers: int, optional
+        The number of decoder layers in the decoder.
+    d_ffn: int, optional
+        The dimension of the feedforward network model hidden layer.
+    dropout: int, optional
+        The dropout value.
+    activation: torch.nn.Module, optional
+        The activation function for Feed-Forward Network layer,
+        e.g., relu or gelu or swish.
+    custom_src_module: torch.nn.Module, optional
+        Module that processes the src features to expected feature dim.
+    custom_tgt_module: torch.nn.Module, optional
+        Module that processes the src features to expected feature dim.
+    positional_encoding: str, optional
+        Type of positional encoding used. e.g. 'fixed_abs_sine' for fixed absolute positional encodings.
+    normalize_before: bool, optional
+        Whether normalization should be applied before or after MHA or FFN in Transformer layers.
+        Defaults to True as this was shown to lead to better performance and training stability.
+    kernel_size: int, optional
+        Kernel size in convolutional layers when Conformer is used.
+    bias: bool, optional
+        Whether to use bias in Conformer convolutional layers.
+    encoder_module: str, optional
+        Choose between Branchformer, Conformer, ConMamba, and Transformer for the encoder.
+    decoder_module: str, optional
+        Choose between Mamba and Transformer for the decoder.
+    conformer_activation: torch.nn.Module, optional
+        Activation module used after Conformer convolutional layers. E.g. Swish, ReLU etc. it has to be a torch Module.
+    branchformer_activation: torch.nn.Module, optional
+        Activation module used within the Branchformer Encoder. E.g. Swish, ReLU etc. it has to be a torch Module.
+    attention_type: str, optional
+        Type of attention layer used in all Transformer or Conformer layers.
+        e.g. regularMHA or RelPosMHA.
+    max_length: int, optional
+        Max length for the target and source sequence in input.
+        Used for positional encodings.
+    causal: bool, optional
+        Whether the encoder should be causal or not (the decoder is always causal).
+        If causal the Conformer convolutional layer is causal.
+    encoder_kdim: int, optional
+        Dimension of the key for the encoder.
+    encoder_vdim: int, optional
+        Dimension of the value for the encoder.
+    decoder_kdim: int, optional
+        Dimension of the key for the decoder.
+    decoder_vdim: int, optional
+        Dimension of the value for the decoder.
+    csgu_linear_units: int, optional
+        Number of neurons in the hidden linear units of the CSGU Module.
+        -> Branchformer
+    gate_activation: torch.nn.Module, optional
+        Activation function used at the gate of the CSGU module.
+        -> Branchformer
+    use_linear_after_conv: bool, optional
+        If True, will apply a linear transformation of size input_size//2.
+        -> Branchformer
+    mamba_config: dict, optional
+        Mamba parameters if encoder_module or decoder_module is Mamba or ConMamba
+    """
+
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        d_ffn=2048,
+        dropout=0.1,
+        activation=nn.ReLU,
+        custom_src_module=None,
+        custom_tgt_module=None,
+        positional_encoding="fixed_abs_sine",
+        normalize_before=True,
+        kernel_size: Optional[int] = 31,
+        bias: Optional[bool] = True,
+        encoder_module: Optional[str] = "transformer",
+        decoder_module: Optional[str] = "transformer",
+        conformer_activation: Optional[nn.Module] = Swish,
+        branchformer_activation: Optional[nn.Module] = nn.GELU,
+        attention_type: Optional[str] = "regularMHA",
+        max_length: Optional[int] = 2500,
+        causal: Optional[bool] = False,
+        encoder_kdim: Optional[int] = None,
+        encoder_vdim: Optional[int] = None,
+        decoder_kdim: Optional[int] = None,
+        decoder_vdim: Optional[int] = None,
+        csgu_linear_units: Optional[int] = 3072,
+        gate_activation: Optional[nn.Module] = nn.Identity,
+        use_linear_after_conv: Optional[bool] = False,
+        mamba_config=None
+    ):
+        super().__init__()
+        self.causal = causal
+        self.attention_type = attention_type
+        self.positional_encoding_type = positional_encoding
+        self.encoder_kdim = encoder_kdim
+        self.encoder_vdim = encoder_vdim
+        self.decoder_kdim = decoder_kdim
+        self.decoder_vdim = decoder_vdim
+
+        assert attention_type in ["regularMHA", "RelPosMHAXL", "hypermixing"]
+        assert positional_encoding in ["fixed_abs_sine", None]
+
+        assert (
+            num_encoder_layers + num_decoder_layers > 0
+        ), "number of encoder layers and number of decoder layers cannot both be 0!"
+
+        if positional_encoding == "fixed_abs_sine":
+            self.positional_encoding = PositionalEncoding(d_model, max_length)
+        elif positional_encoding is None:
+            pass
+            # no positional encodings
+
+        # overrides any other pos_embedding
+        if attention_type == "RelPosMHAXL":
+            self.positional_encoding = RelPosEncXL(d_model)
+            self.positional_encoding_decoder = PositionalEncoding(
+                d_model, max_length
+            )
+
+        # initialize the encoder
+        if num_encoder_layers > 0:
+            if custom_src_module is not None:
+                self.custom_src_module = custom_src_module(d_model)
+            if encoder_module == "transformer":
+                self.encoder = TransformerEncoder(
+                    nhead=nhead,
+                    num_layers=num_encoder_layers,
+                    d_ffn=d_ffn,
+                    d_model=d_model,
+                    dropout=dropout,
+                    activation=activation,
+                    normalize_before=normalize_before,
+                    causal=self.causal,
+                    attention_type=self.attention_type,
+                    kdim=self.encoder_kdim,
+                    vdim=self.encoder_vdim,
+                )
+            elif encoder_module == "conformer":
+                self.encoder = ConformerEncoder(
+                    nhead=nhead,
+                    num_layers=num_encoder_layers,
+                    d_ffn=d_ffn,
+                    d_model=d_model,
+                    dropout=dropout,
+                    activation=conformer_activation,
+                    kernel_size=kernel_size,
+                    bias=bias,
+                    causal=self.causal,
+                    attention_type=self.attention_type,
+                )
+                assert (
+                    normalize_before
+                ), "normalize_before must be True for Conformer"
+
+                assert (
+                    conformer_activation is not None
+                ), "conformer_activation must not be None"
+            elif encoder_module == "branchformer":
+                self.encoder = BranchformerEncoder(
+                    nhead=nhead,
+                    num_layers=num_encoder_layers,
+                    d_model=d_model,
+                    dropout=dropout,
+                    activation=branchformer_activation,
+                    kernel_size=kernel_size,
+                    attention_type=self.attention_type,
+                    csgu_linear_units=csgu_linear_units,
+                    gate_activation=gate_activation,
+                    use_linear_after_conv=use_linear_after_conv,
+                )
+            elif encoder_module == "conmamba":
+                self.encoder = ConmambaEncoder(
+                    num_layers=num_encoder_layers,
+                    d_model=d_model,
+                    d_ffn=d_ffn,
+                    dropout=dropout,
+                    activation=branchformer_activation,
+                    kernel_size=kernel_size,
+                    bias=bias,
+                    causal=self.causal,
+                    mamba_config=mamba_config
+                )
+                assert (
+                    normalize_before
+                ), "normalize_before must be True for Conmamba"
+
+                assert (
+                    conformer_activation is not None
+                ), "conformer_activation must not be None"
+
+        # initialize the decoder
+        if num_decoder_layers > 0:
+            if custom_tgt_module is not None:
+                self.custom_tgt_module = custom_tgt_module(d_model)
+            if decoder_module == 'transformer':
+                self.decoder = TransformerDecoder(
+                    num_layers=num_decoder_layers,
+                    nhead=nhead,
+                    d_ffn=d_ffn,
+                    d_model=d_model,
+                    dropout=dropout,
+                    activation=activation,
+                    normalize_before=normalize_before,
+                    causal=True,
+                    attention_type="regularMHA",  # always use regular attention in decoder
+                    kdim=self.decoder_kdim,
+                    vdim=self.decoder_vdim,
+                )
+            elif decoder_module in ['mamba']:
+                self.decoder = MambaDecoder(
+                    num_layers=num_decoder_layers,
+                    d_ffn=d_ffn,
+                    d_model=d_model,
+                    activation=activation,
+                    dropout=dropout,
+                    normalize_before=normalize_before,
+                    mamba_config=mamba_config
+                )
+            else:
+                raise NotImplementedError(decoder_module)
+
+    def forward(self, **kwags):
+        """Users should modify this function according to their own tasks."""
+        raise NotImplementedError
+
+
+class PositionalEncoding(nn.Module):
+    """This class implements the absolute sinusoidal positional encoding function.
+    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
+    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
+
+    Arguments
+    ---------
+    input_size: int
+        Embedding dimension.
+    max_len : int, optional
+        Max length of the input sequences (default 2500).
+
+    Example
+    -------
+    >>> a = torch.rand((8, 120, 512))
+    >>> enc = PositionalEncoding(input_size=a.shape[-1])
+    >>> b = enc(a)
+    >>> b.shape
+    torch.Size([1, 120, 512])
+    """
+
+    def __init__(self, input_size, max_len=2500):
+        super().__init__()
+        if input_size % 2 != 0:
+            raise ValueError(
+                f"Cannot use sin/cos positional encoding with odd channels (got channels={input_size})"
+            )
+        self.max_len = max_len
+        pe = torch.zeros(self.max_len, input_size, requires_grad=False)
+        positions = torch.arange(0, self.max_len).unsqueeze(1).float()
+        denominator = torch.exp(
+            torch.arange(0, input_size, 2).float()
+            * -(math.log(10000.0) / input_size)
+        )
+
+        pe[:, 0::2] = torch.sin(positions * denominator)
+        pe[:, 1::2] = torch.cos(positions * denominator)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input feature shape (batch, time, fea)
+
+        Returns
+        -------
+        The positional encoding.
+        """
+        return self.pe[:, : x.size(1)].clone().detach()
+
+
+class TransformerEncoderLayer(nn.Module):
+    """This is an implementation of self-attention encoder layer.
+
+    Arguments
+    ---------
+    d_ffn: int, optional
+        The dimension of the feedforward network model hidden layer.
+    nhead: int
+        The number of heads in the multi-head attention models (default=8).
+    d_model: int
+        The number of expected features in the encoder/decoder inputs (default=512).
+    kdim: int, optional
+        Dimension of the key.
+    vdim: int, optional
+        Dimension of the value.
+    dropout: int, optional
+        The dropout value.
+    activation: torch.nn.Module, optional
+        The activation function for Feed-Forward Network layer,
+        e.g., relu or gelu or swish.
+    normalize_before: bool, optional
+        Whether normalization should be applied before or after MHA or FFN in Transformer layers.
+        Defaults to True as this was shown to lead to better performance and training stability.
+    attention_type: str, optional
+        Type of attention layer used in all Transformer or Conformer layers.
+        e.g. regularMHA or RelPosMHA.
+    ffn_type: str
+        type of ffn: regularFFN/1dcnn
+    ffn_cnn_kernel_size_list: list of int
+        kernel size of 2 1d-convs if ffn_type is 1dcnn
+    causal: bool, optional
+        Whether the encoder should be causal or not (the decoder is always causal).
+        If causal the Conformer convolutional layer is causal.
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> net = TransformerEncoderLayer(512, 8, d_model=512)
+    >>> output = net(x)
+    >>> output[0].shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        d_ffn,
+        nhead,
+        d_model,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        activation=nn.ReLU,
+        normalize_before=False,
+        attention_type="regularMHA",
+        ffn_type="regularFFN",
+        ffn_cnn_kernel_size_list=[3, 3],
+        causal=False,
+    ):
+        super().__init__()
+
+        if attention_type == "regularMHA":
+            self.self_att = sb.nnet.attention.MultiheadAttention(
+                nhead=nhead,
+                d_model=d_model,
+                dropout=dropout,
+                kdim=kdim,
+                vdim=vdim,
+            )
+
+        elif attention_type == "RelPosMHAXL":
+            self.self_att = sb.nnet.attention.RelPosMHAXL(
+                d_model, nhead, dropout, mask_pos_future=causal
+            )
+        elif attention_type == "hypermixing":
+            self.self_att = sb.nnet.hypermixing.HyperMixing(
+                input_output_dim=d_model,
+                hypernet_size=d_ffn,
+                tied=False,
+                num_heads=nhead,
+                fix_tm_hidden_size=False,
+            )
+
+        if ffn_type == "regularFFN":
+            self.pos_ffn = sb.nnet.attention.PositionalwiseFeedForward(
+                d_ffn=d_ffn,
+                input_size=d_model,
+                dropout=dropout,
+                activation=activation,
+            )
+        elif ffn_type == "1dcnn":
+            self.pos_ffn = nn.Sequential(
+                Conv1d(
+                    in_channels=d_model,
+                    out_channels=d_ffn,
+                    kernel_size=ffn_cnn_kernel_size_list[0],
+                    padding="causal" if causal else "same",
+                ),
+                nn.ReLU(),
+                Conv1d(
+                    in_channels=d_ffn,
+                    out_channels=d_model,
+                    kernel_size=ffn_cnn_kernel_size_list[1],
+                    padding="causal" if causal else "same",
+                ),
+            )
+
+        self.norm1 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.norm2 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+
+        self.normalize_before = normalize_before
+        self.pos_ffn_type = ffn_type
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: Optional[torch.Tensor] = None,
+    ):
+        """
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence to the encoder layer.
+        src_mask : torch.Tensor
+            The mask for the src query for each example in the batch.
+        src_key_padding_mask : torch.Tensor, optional
+            The mask for the src keys for each example in the batch.
+        pos_embs: torch.Tensor, optional
+            The positional embeddings tensor.
+
+        Returns
+        -------
+        output : torch.Tensor
+            The output of the transformer encoder layer.
+        """
+
+        if self.normalize_before:
+            src1 = self.norm1(src)
+        else:
+            src1 = src
+
+        output, self_attn = self.self_att(
+            src1,
+            src1,
+            src1,
+            attn_mask=src_mask,
+            key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs,
+        )
+
+        # add & norm
+        src = src + self.dropout1(output)
+        if not self.normalize_before:
+            src = self.norm1(src)
+
+        if self.normalize_before:
+            src1 = self.norm2(src)
+        else:
+            src1 = src
+        output = self.pos_ffn(src1)
+
+        # add & norm
+        output = src + self.dropout2(output)
+        if not self.normalize_before:
+            output = self.norm2(output)
+        return output, self_attn
+
+
+class TransformerEncoder(nn.Module):
+    """This class implements the transformer encoder.
+
+    Arguments
+    ---------
+    num_layers : int
+        Number of transformer layers to include.
+    nhead : int
+        Number of attention heads.
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    input_shape : tuple
+        Expected shape of the input.
+    d_model : int
+        The dimension of the input embedding.
+    kdim : int
+        Dimension for key (Optional).
+    vdim : int
+        Dimension for value (Optional).
+    dropout : float
+        Dropout for the encoder (Optional).
+    activation: torch.nn.Module, optional
+        The activation function for Feed-Forward Network layer,
+        e.g., relu or gelu or swish.
+    normalize_before: bool, optional
+        Whether normalization should be applied before or after MHA or FFN in Transformer layers.
+        Defaults to True as this was shown to lead to better performance and training stability.
+    causal: bool, optional
+        Whether the encoder should be causal or not (the decoder is always causal).
+        If causal the Conformer convolutional layer is causal.
+    layerdrop_prob: float
+        The probability to drop an entire layer
+    attention_type: str, optional
+        Type of attention layer used in all Transformer or Conformer layers.
+        e.g. regularMHA or RelPosMHA.
+    ffn_type: str
+        type of ffn: regularFFN/1dcnn
+    ffn_cnn_kernel_size_list: list of int
+        conv kernel size of 2 1d-convs if ffn_type is 1dcnn
+
+    Example
+    -------
+    >>> import torch
+    >>> x = torch.rand((8, 60, 512))
+    >>> net = TransformerEncoder(1, 8, 512, d_model=512)
+    >>> output, _ = net(x)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        nhead,
+        d_ffn,
+        input_shape=None,
+        d_model=None,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        activation=nn.ReLU,
+        normalize_before=False,
+        causal=False,
+        layerdrop_prob=0.0,
+        attention_type="regularMHA",
+        ffn_type="regularFFN",
+        ffn_cnn_kernel_size_list=[3, 3],
+    ):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(
+            [
+                TransformerEncoderLayer(
+                    d_ffn=d_ffn,
+                    nhead=nhead,
+                    d_model=d_model,
+                    kdim=kdim,
+                    vdim=vdim,
+                    dropout=dropout,
+                    activation=activation,
+                    normalize_before=normalize_before,
+                    causal=causal,
+                    attention_type=attention_type,
+                    ffn_type=ffn_type,
+                    ffn_cnn_kernel_size_list=ffn_cnn_kernel_size_list,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.norm = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.layerdrop_prob = layerdrop_prob
+        self.rng = np.random.default_rng()
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos_embs: Optional[torch.Tensor] = None,
+        dynchunktrain_config=None,
+    ):
+        """
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence to the encoder layer (required).
+        src_mask : torch.Tensor
+            The mask for the src sequence (optional).
+        src_key_padding_mask : torch.Tensor
+            The mask for the src keys per batch (optional).
+        pos_embs : torch.Tensor
+            The positional embedding tensor
+        dynchunktrain_config : config
+            Not supported for this encoder.
+
+        Returns
+        -------
+        output : torch.Tensor
+            The output of the transformer.
+        attention_lst : list
+            The attention values.
+        """
+        assert (
+            dynchunktrain_config is None
+        ), "Dynamic Chunk Training unsupported for this encoder"
+
+        output = src
+        if self.layerdrop_prob > 0.0:
+            keep_probs = self.rng.random(len(self.layers))
+        else:
+            keep_probs = None
+        attention_lst = []
+        for i, enc_layer in enumerate(self.layers):
+            if (
+                not self.training
+                or self.layerdrop_prob == 0.0
+                or keep_probs[i] > self.layerdrop_prob
+            ):
+                output, attention = enc_layer(
+                    output,
+                    src_mask=src_mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    pos_embs=pos_embs,
+                )
+
+                attention_lst.append(attention)
+        output = self.norm(output)
+        return output, attention_lst
+
+
+class TransformerDecoderLayer(nn.Module):
+    """This class implements the self-attention decoder layer.
+
+    Arguments
+    ---------
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    nhead : int
+        Number of attention heads.
+    d_model : int
+        Dimension of the model.
+    kdim : int
+        Dimension for key (optional).
+    vdim : int
+        Dimension for value (optional).
+    dropout : float
+        Dropout for the decoder (optional).
+    activation : Callable
+        Function to use between layers, default nn.ReLU
+    normalize_before : bool
+        Whether to normalize before layers.
+    attention_type : str
+        Type of attention to use, "regularMHA" or "RelPosMHAXL"
+    causal : bool
+        Whether to mask future positions.
+
+    Example
+    -------
+    >>> src = torch.rand((8, 60, 512))
+    >>> tgt = torch.rand((8, 60, 512))
+    >>> net = TransformerDecoderLayer(1024, 8, d_model=512)
+    >>> output, self_attn, multihead_attn = net(src, tgt)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        d_ffn,
+        nhead,
+        d_model,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        activation=nn.ReLU,
+        normalize_before=False,
+        attention_type="regularMHA",
+        causal=None,
+    ):
+        super().__init__()
+        self.nhead = nhead
+
+        if attention_type == "regularMHA":
+            self.self_attn = sb.nnet.attention.MultiheadAttention(
+                nhead=nhead,
+                d_model=d_model,
+                kdim=kdim,
+                vdim=vdim,
+                dropout=dropout,
+            )
+            self.multihead_attn = sb.nnet.attention.MultiheadAttention(
+                nhead=nhead,
+                d_model=d_model,
+                kdim=kdim,
+                vdim=vdim,
+                dropout=dropout,
+            )
+
+        elif attention_type == "RelPosMHAXL":
+            self.self_attn = sb.nnet.attention.RelPosMHAXL(
+                d_model, nhead, dropout, mask_pos_future=causal
+            )
+            self.multihead_attn = sb.nnet.attention.RelPosMHAXL(
+                d_model, nhead, dropout, mask_pos_future=causal
+            )
+
+        self.pos_ffn = sb.nnet.attention.PositionalwiseFeedForward(
+            d_ffn=d_ffn,
+            input_size=d_model,
+            dropout=dropout,
+            activation=activation,
+        )
+
+        # normalization layers
+        self.norm1 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.norm2 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.norm3 = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+        self.dropout3 = torch.nn.Dropout(dropout)
+
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ----------
+        tgt: torch.Tensor
+            The sequence to the decoder layer (required).
+        memory: torch.Tensor
+            The sequence from the last layer of the encoder (required).
+        tgt_mask: torch.Tensor
+            The mask for the tgt sequence (optional).
+        memory_mask: torch.Tensor
+            The mask for the memory sequence (optional).
+        tgt_key_padding_mask: torch.Tensor
+            The mask for the tgt keys per batch (optional).
+        memory_key_padding_mask: torch.Tensor
+            The mask for the memory keys per batch (optional).
+        pos_embs_tgt: torch.Tensor
+            The positional embeddings for the target (optional).
+        pos_embs_src: torch.Tensor
+            The positional embeddings for the source (optional).
+        """
+        if self.normalize_before:
+            tgt1 = self.norm1(tgt)
+        else:
+            tgt1 = tgt
+
+        # self-attention over the target sequence
+        tgt2, self_attn = self.self_attn(
+            query=tgt1,
+            key=tgt1,
+            value=tgt1,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask,
+            pos_embs=pos_embs_tgt,
+        )
+
+        # add & norm
+        tgt = tgt + self.dropout1(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm1(tgt)
+
+        if self.normalize_before:
+            tgt1 = self.norm2(tgt)
+        else:
+            tgt1 = tgt
+
+        # multi-head attention over the target sequence and encoder states
+
+        tgt2, multihead_attention = self.multihead_attn(
+            query=tgt1,
+            key=memory,
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+            pos_embs=pos_embs_src,
+        )
+
+        # add & norm
+        tgt = tgt + self.dropout2(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm2(tgt)
+
+        if self.normalize_before:
+            tgt1 = self.norm3(tgt)
+        else:
+            tgt1 = tgt
+
+        tgt2 = self.pos_ffn(tgt1)
+
+        # add & norm
+        tgt = tgt + self.dropout3(tgt2)
+        if not self.normalize_before:
+            tgt = self.norm3(tgt)
+
+        return tgt, self_attn, multihead_attention
+
+
+class TransformerDecoder(nn.Module):
+    """This class implements the Transformer decoder.
+
+    Arguments
+    ---------
+    num_layers : int
+        Number of transformer layers for the decoder.
+    nhead : int
+        Number of attention heads.
+    d_ffn : int
+        Hidden size of self-attention Feed Forward layer.
+    d_model : int
+        Dimension of the model.
+    kdim : int, optional
+        Dimension for key (Optional).
+    vdim : int, optional
+        Dimension for value (Optional).
+    dropout : float, optional
+        Dropout for the decoder (Optional).
+    activation : Callable
+        The function to apply between layers, default nn.ReLU
+    normalize_before : bool
+        Whether to normalize before layers.
+    causal : bool
+        Whether to allow future information in decoding.
+    attention_type : str
+        Type of attention to use, "regularMHA" or "RelPosMHAXL"
+
+    Example
+    -------
+    >>> src = torch.rand((8, 60, 512))
+    >>> tgt = torch.rand((8, 60, 512))
+    >>> net = TransformerDecoder(1, 8, 1024, d_model=512)
+    >>> output, _, _ = net(src, tgt)
+    >>> output.shape
+    torch.Size([8, 60, 512])
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        nhead,
+        d_ffn,
+        d_model,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        activation=nn.ReLU,
+        normalize_before=False,
+        causal=False,
+        attention_type="regularMHA",
+    ):
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                TransformerDecoderLayer(
+                    d_ffn=d_ffn,
+                    nhead=nhead,
+                    d_model=d_model,
+                    kdim=kdim,
+                    vdim=vdim,
+                    dropout=dropout,
+                    activation=activation,
+                    normalize_before=normalize_before,
+                    causal=causal,
+                    attention_type=attention_type,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = sb.nnet.normalization.LayerNorm(d_model, eps=1e-6)
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+        pos_embs_tgt=None,
+        pos_embs_src=None,
+    ):
+        """
+        Arguments
+        ----------
+        tgt : torch.Tensor
+            The sequence to the decoder layer (required).
+        memory : torch.Tensor
+            The sequence from the last layer of the encoder (required).
+        tgt_mask : torch.Tensor
+            The mask for the tgt sequence (optional).
+        memory_mask : torch.Tensor
+            The mask for the memory sequence (optional).
+        tgt_key_padding_mask : torch.Tensor
+            The mask for the tgt keys per batch (optional).
+        memory_key_padding_mask : torch.Tensor
+            The mask for the memory keys per batch (optional).
+        pos_embs_tgt : torch.Tensor
+            The positional embeddings for the target (optional).
+        pos_embs_src : torch.Tensor
+            The positional embeddings for the source (optional).
+        """
+        output = tgt
+        self_attns, multihead_attns = [], []
+        for dec_layer in self.layers:
+            output, self_attn, multihead_attn = dec_layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                pos_embs_tgt=pos_embs_tgt,
+                pos_embs_src=pos_embs_src,
+            )
+            self_attns.append(self_attn)
+            multihead_attns.append(multihead_attn)
+        output = self.norm(output)
+
+        return output, self_attns, multihead_attns
+
+
+class NormalizedEmbedding(nn.Module):
+    """This class implements the normalized embedding layer for the transformer.
+    Since the dot product of the self-attention is always normalized by sqrt(d_model)
+    and the final linear projection for prediction shares weight with the embedding layer,
+    we multiply the output of the embedding by sqrt(d_model).
+
+    Arguments
+    ---------
+    d_model: int
+        The number of expected features in the encoder/decoder inputs (default=512).
+    vocab: int
+        The vocab size.
+
+    Example
+    -------
+    >>> emb = NormalizedEmbedding(512, 1000)
+    >>> trg = torch.randint(0, 999, (8, 50))
+    >>> emb_fea = emb(trg)
+    """
+
+    def __init__(self, d_model, vocab):
+        super().__init__()
+        self.emb = sb.nnet.embedding.Embedding(
+            num_embeddings=vocab, embedding_dim=d_model, blank_id=0
+        )
+        self.d_model = d_model
+
+    def forward(self, x):
+        """Processes the input tensor x and returns an output tensor."""
+        return self.emb(x) * math.sqrt(self.d_model)
+
+
+def get_key_padding_mask(padded_input, pad_idx):
+    """Creates a binary mask to prevent attention to padded locations.
+    We suggest using ``get_mask_from_lengths`` instead of this function.
+
+    Arguments
+    ---------
+    padded_input: torch.Tensor
+        Padded input.
+    pad_idx: int
+        idx for padding element.
+
+    Returns
+    -------
+    key_padded_mask: torch.Tensor
+        Binary mask to prevent attention to padding.
+
+    Example
+    -------
+    >>> a = torch.LongTensor([[1,1,0], [2,3,0], [4,5,0]])
+    >>> get_key_padding_mask(a, pad_idx=0)
+    tensor([[False, False,  True],
+            [False, False,  True],
+            [False, False,  True]])
+    """
+    if len(padded_input.shape) == 4:
+        bz, time, ch1, ch2 = padded_input.shape
+        padded_input = padded_input.reshape(bz, time, ch1 * ch2)
+
+    key_padded_mask = padded_input.eq(pad_idx).to(padded_input.device)
+
+    # if the input is more than 2d, mask the locations where they are silence
+    # across all channels
+    if len(padded_input.shape) > 2:
+        key_padded_mask = key_padded_mask.float().prod(dim=-1).bool()
+        return key_padded_mask.detach()
+
+    return key_padded_mask.detach()
+
+
+def get_lookahead_mask(padded_input):
+    """Creates a binary mask for each sequence which masks future frames.
+
+    Arguments
+    ---------
+    padded_input: torch.Tensor
+        Padded input tensor.
+
+    Returns
+    -------
+    mask : torch.Tensor
+        Binary mask for masking future frames.
+
+    Example
+    -------
+    >>> a = torch.LongTensor([[1,1,0], [2,3,0], [4,5,0]])
+    >>> get_lookahead_mask(a)
+    tensor([[0., -inf, -inf],
+            [0., 0., -inf],
+            [0., 0., 0.]])
+    """
+    seq_len = padded_input.shape[1]
+    mask = (
+        torch.triu(torch.ones((seq_len, seq_len), device=padded_input.device))
+        == 1
+    ).transpose(0, 1)
+    mask = (
+        mask.float()
+        .masked_fill(mask == 0, float("-inf"))
+        .masked_fill(mask == 1, float(0.0))
+    )
+    return mask.detach().to(padded_input.device)
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    """Creates a binary mask from sequence lengths
+
+    Arguments
+    ---------
+    lengths: torch.Tensor
+        A tensor of sequence lengths
+    max_len: int (Optional)
+        Maximum sequence length, defaults to None.
+
+    Returns
+    -------
+    mask: torch.Tensor
+        the mask where padded elements are set to True.
+        Then one can use tensor.masked_fill_(mask, 0) for the masking.
+
+    Example
+    -------
+    >>> lengths = torch.tensor([3, 2, 4])
+    >>> get_mask_from_lengths(lengths)
+    tensor([[False, False, False,  True],
+            [False, False,  True,  True],
+            [False, False, False, False]])
+    """
+    if max_len is None:
+        max_len = torch.max(lengths).item()
+    seq_range = torch.arange(
+        max_len, device=lengths.device, dtype=lengths.dtype
+    )
+    return ~(seq_range.unsqueeze(0) < lengths.unsqueeze(1))

modules/TransformerASR.py

@@ -0,0 +1,682 @@
+"""Added ConMamba and Mamba
+
+Authors
+* Xilin Jiang 2024
+"""
+
+"""Transformer for ASR in the SpeechBrain style.
+
+Authors
+* Jianyuan Zhong 2020
+* Titouan Parcollet 2024
+* Luca Della Libera 2024
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch  # noqa 42
+from torch import nn
+
+from speechbrain.dataio.dataio import length_to_mask
+from modules.Transformer import (
+    NormalizedEmbedding,
+    TransformerInterface,
+    get_key_padding_mask,
+    get_lookahead_mask,
+)
+from speechbrain.nnet.activations import Swish
+from speechbrain.nnet.containers import ModuleList
+from speechbrain.nnet.linear import Linear
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+
+
+@dataclass
+class TransformerASRStreamingContext:
+    """Streaming metadata and state for a `TransformerASR` instance."""
+
+    dynchunktrain_config: DynChunkTrainConfig
+    """Dynamic Chunk Training configuration holding chunk size and context size
+    information."""
+
+    encoder_context: Any
+    """Opaque encoder context information. It is constructed by the encoder's
+    `make_streaming_context` method and is passed to the encoder when using
+    `encode_streaming`.
+    """
+
+
+def make_transformer_src_mask(
+    src: torch.Tensor,
+    causal: bool = False,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+) -> Optional[torch.Tensor]:
+    """Prepare the source transformer mask that restricts which frames can
+    attend to which frames depending on causal or other simple restricted
+    attention methods.
+
+    Arguments
+    ---------
+    src: torch.Tensor
+        The source tensor to build a mask from. The contents of the tensor are
+        not actually used currently; only its shape and other metadata (e.g.
+        device).
+    causal: bool
+        Whether strict causality shall be used. Frames will not be able to
+        attend to any future frame.
+    dynchunktrain_config: DynChunkTrainConfig, optional
+        Dynamic Chunk Training configuration. This implements a simple form of
+        chunkwise attention. Incompatible with `causal`.
+
+    Returns
+    -------
+    torch.Tensor
+        A boolean mask Tensor of shape (timesteps, timesteps).
+    """
+    if causal:
+        assert dynchunktrain_config is None
+        return get_lookahead_mask(src)
+
+    if dynchunktrain_config is None:
+        return
+
+    # The following is not really the sole source used to implement this,
+    # but it helps introduce the concept.
+    # ref: Unified Streaming and Non-streaming Two-pass End-to-end Model for Speech Recognition
+    # https://arxiv.org/pdf/2012.05481.pdf
+    timesteps = src.size(1)
+
+    # Mask the future at the right of each chunk
+    chunk_size = dynchunktrain_config.chunk_size
+    num_chunks = timesteps // chunk_size
+    timestep_idx = torch.arange(timesteps, device=src.device)
+    mask_idx = torch.arange(
+        chunk_size, chunk_size * (num_chunks + 2), chunk_size, device=src.device
+    ).repeat_interleave(chunk_size)[:timesteps]
+    src_mask = timestep_idx[None] >= mask_idx[:, None]
+
+    # Mask the past at the left of each chunk (accounting for left context)
+    # only relevant if using left context
+    if not dynchunktrain_config.is_infinite_left_context():
+        num_left_chunks = dynchunktrain_config.left_context_size
+        mask_idx -= chunk_size * (num_left_chunks + 1)
+        src_mask += timestep_idx[None] < mask_idx[:, None]
+
+    return src_mask
+
+
+def make_transformer_src_tgt_masks(
+    src,
+    tgt=None,
+    wav_len=None,
+    pad_idx=0,
+    causal: bool = False,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+):
+    """This function generates masks for training the transformer model,
+    opinionated for an ASR context with encoding masks and, optionally, decoding
+    masks (if specifying `tgt`).
+
+    Arguments
+    ---------
+    src : torch.Tensor
+        The sequence to the encoder (required).
+    tgt : torch.Tensor
+        The sequence to the decoder.
+    wav_len : torch.Tensor
+        The lengths of the inputs.
+    pad_idx : int
+        The index for <pad> token (default=0).
+    causal: bool
+        Whether strict causality shall be used. See `make_asr_src_mask`
+    dynchunktrain_config: DynChunkTrainConfig, optional
+        Dynamic Chunk Training configuration. See `make_asr_src_mask`
+
+    Returns
+    -------
+    src_key_padding_mask : torch.Tensor
+        Key padding mask for ignoring padding
+    tgt_key_padding_mask : torch.Tensor
+        Key padding mask for ignoring padding
+    src_mask : torch.Tensor
+        Mask for ignoring invalid (e.g. future) timesteps
+    tgt_mask : torch.Tensor
+        Mask for ignoring invalid (e.g. future) timesteps
+    """
+    src_key_padding_mask = None
+
+    # mask out audio beyond the length of audio for each batch
+    if wav_len is not None:
+        abs_len = torch.round(wav_len * src.shape[1])
+        src_key_padding_mask = ~length_to_mask(abs_len).bool()
+
+    # mask out the source
+    src_mask = make_transformer_src_mask(
+        src, causal=causal, dynchunktrain_config=dynchunktrain_config
+    )
+
+    # If no decoder in the transformer...
+    if tgt is not None:
+        tgt_key_padding_mask = get_key_padding_mask(tgt, pad_idx=pad_idx)
+        tgt_mask = get_lookahead_mask(tgt)
+    else:
+        tgt_key_padding_mask = None
+        tgt_mask = None
+
+    return src_key_padding_mask, tgt_key_padding_mask, src_mask, tgt_mask
+
+
+class TransformerASR(TransformerInterface):
+    """This is an implementation of transformer model for ASR.
+
+    The architecture is based on the paper "Attention Is All You Need":
+    https://arxiv.org/pdf/1706.03762.pdf
+
+    Arguments
+    ---------
+    tgt_vocab: int
+        Size of vocabulary.
+    input_size: int
+        Input feature size.
+    d_model : int, optional
+        Embedding dimension size.
+        (default=512).
+    nhead : int, optional
+        The number of heads in the multi-head attention models (default=8).
+    num_encoder_layers : int, optional
+        The number of sub-encoder-layers in the encoder (default=6).
+    num_decoder_layers : int, optional
+        The number of sub-decoder-layers in the decoder (default=6).
+    d_ffn : int, optional
+        The dimension of the feedforward network model (default=2048).
+    dropout : int, optional
+        The dropout value (default=0.1).
+    activation : torch.nn.Module, optional
+        The activation function of FFN layers.
+        Recommended: relu or gelu (default=relu).
+    positional_encoding: str, optional
+        Type of positional encoding used. e.g. 'fixed_abs_sine' for fixed absolute positional encodings.
+    normalize_before: bool, optional
+        Whether normalization should be applied before or after MHA or FFN in Transformer layers.
+        Defaults to True as this was shown to lead to better performance and training stability.
+    kernel_size: int, optional
+        Kernel size in convolutional layers when Conformer is used.
+    bias: bool, optional
+        Whether to use bias in Conformer convolutional layers.
+    encoder_module: str, optional
+        Choose between Branchformer, Conformer, ConMamba, and Transformer for the encoder.
+    decoder_module: str, optional
+        Choose between Mamba and Transformer for the decoder.
+    decoder_module: str, optional
+        Choose between Transformer and Mamba for the decoder. 
+    conformer_activation: torch.nn.Module, optional
+        Activation module used after Conformer convolutional layers. E.g. Swish, ReLU etc. it has to be a torch Module.
+    branchformer_activation: torch.nn.Module, optional
+        Activation module used within the Branchformer Encoder. E.g. Swish, ReLU etc. it has to be a torch Module.
+    attention_type: str, optional
+        Type of attention layer used in all Transformer or Conformer layers.
+        e.g. regularMHA or RelPosMHA.
+    max_length: int, optional
+        Max length for the target and source sequence in input.
+        Used for positional encodings.
+    causal: bool, optional
+        Whether the encoder should be causal or not (the decoder is always causal).
+        If causal the Conformer convolutional layer is causal.
+    csgu_linear_units: int, optional
+        Number of neurons in the hidden linear units of the CSGU Module.
+        -> Branchformer
+    gate_activation: torch.nn.Module, optional
+        Activation function used at the gate of the CSGU module.
+        -> Branchformer
+    use_linear_after_conv: bool, optional
+        If True, will apply a linear transformation of size input_size//2.
+        -> Branchformer
+    mamba_config: dict, optional
+        Mamba parameters if encoder_module or decoder_module is Mamba or ConMamba
+
+    Example
+    -------
+    >>> src = torch.rand([8, 120, 512])
+    >>> tgt = torch.randint(0, 720, [8, 120])
+    >>> net = TransformerASR(
+    ...     720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU
+    ... )
+    >>> enc_out, dec_out = net.forward(src, tgt)
+    >>> enc_out.shape
+    torch.Size([8, 120, 512])
+    >>> dec_out.shape
+    torch.Size([8, 120, 512])
+    """
+
+    def __init__(
+        self,
+        tgt_vocab,
+        input_size,
+        d_model=512,
+        nhead=8,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        d_ffn=2048,
+        dropout=0.1,
+        activation=nn.ReLU,
+        positional_encoding="fixed_abs_sine",
+        normalize_before=False,
+        kernel_size: Optional[int] = 31,
+        bias: Optional[bool] = True,
+        encoder_module: Optional[str] = "transformer",
+        decoder_module: Optional[str] = "transformer",
+        conformer_activation: Optional[nn.Module] = Swish,
+        branchformer_activation: Optional[nn.Module] = nn.GELU,
+        attention_type: Optional[str] = "regularMHA",
+        max_length: Optional[int] = 2500,
+        causal: Optional[bool] = True,
+        csgu_linear_units: Optional[int] = 3072,
+        gate_activation: Optional[nn.Module] = nn.Identity,
+        use_linear_after_conv: Optional[bool] = False,
+        mamba_config=None
+    ):
+        super().__init__(
+            d_model=d_model,
+            nhead=nhead,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            d_ffn=d_ffn,
+            dropout=dropout,
+            activation=activation,
+            positional_encoding=positional_encoding,
+            normalize_before=normalize_before,
+            kernel_size=kernel_size,
+            bias=bias,
+            encoder_module=encoder_module,
+            decoder_module=decoder_module,
+            conformer_activation=conformer_activation,
+            branchformer_activation=branchformer_activation,
+            attention_type=attention_type,
+            max_length=max_length,
+            causal=causal,
+            csgu_linear_units=csgu_linear_units,
+            gate_activation=gate_activation,
+            use_linear_after_conv=use_linear_after_conv,
+            mamba_config=mamba_config
+        )
+
+        self.custom_src_module = ModuleList(
+            Linear(
+                input_size=input_size,
+                n_neurons=d_model,
+                bias=True,
+                combine_dims=False,
+            ),
+            torch.nn.Dropout(dropout),
+        )
+
+        self.num_decoder_layers = num_decoder_layers
+        if num_decoder_layers > 0:
+            self.custom_tgt_module = ModuleList(
+                NormalizedEmbedding(d_model, tgt_vocab)
+            )
+
+        # reset parameters using xavier_normal_
+        self._init_params()
+
+    def forward(self, src, tgt, wav_len=None, pad_idx=0):
+        """
+        Arguments
+        ----------
+        src : torch.Tensor
+            The sequence to the encoder.
+        tgt : torch.Tensor
+            The sequence to the decoder.
+        wav_len: torch.Tensor, optional
+            Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
+        pad_idx : int, optional
+            The index for <pad> token (default=0).
+        """
+
+        # reshape the src vector to [Batch, Time, Fea] is a 4d vector is given
+        if src.ndim == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        (
+            src_key_padding_mask,
+            tgt_key_padding_mask,
+            src_mask,
+            tgt_mask,
+        ) = make_transformer_src_tgt_masks(
+            src, tgt, wav_len, causal=self.causal, pad_idx=pad_idx
+        )
+        
+        src = self.custom_src_module(src)
+        # add pos encoding to queries if are sinusoidal ones else
+        if self.attention_type == "hypermixing":
+            pos_embs_encoder = None
+        elif self.attention_type == "RelPosMHAXL":
+            pos_embs_encoder = self.positional_encoding(src)
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(src)  # add the encodings here
+            pos_embs_encoder = None
+
+        encoder_out, _ = self.encoder(
+            src=src,
+            src_mask=src_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs_encoder,
+        )
+
+        if self.num_decoder_layers > 0:
+            tgt = self.custom_tgt_module(tgt)
+
+            if self.attention_type == "RelPosMHAXL":
+                tgt = tgt + self.positional_encoding_decoder(tgt)
+                pos_embs_encoder = None  # self.positional_encoding(src)
+                pos_embs_target = None
+            elif (
+                self.positional_encoding_type == "fixed_abs_sine"
+                or self.attention_type == "hypermixing"
+            ):
+                tgt = tgt + self.positional_encoding(tgt)
+                pos_embs_target = None
+                pos_embs_encoder = None
+
+            decoder_out, _, _ = self.decoder(
+                tgt=tgt,
+                memory=encoder_out,
+                memory_mask=None,
+                tgt_mask=tgt_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=src_key_padding_mask,
+                pos_embs_tgt=pos_embs_target,
+                pos_embs_src=pos_embs_encoder,
+            )
+
+        else:
+            decoder_out = None
+
+        return encoder_out, decoder_out
+
+    @torch.no_grad()
+    def decode(self, tgt, encoder_out, enc_len=None):
+        """This method implements a decoding step for the transformer model.
+
+        Arguments
+        ---------
+        tgt : torch.Tensor
+            The sequence to the decoder.
+        encoder_out : torch.Tensor
+            Hidden output of the encoder.
+        enc_len : torch.LongTensor
+            The actual length of encoder states.
+
+        Returns
+        -------
+        prediction
+        """
+        tgt_mask = get_lookahead_mask(tgt)
+        src_key_padding_mask = None
+        if enc_len is not None:
+            src_key_padding_mask = (1 - length_to_mask(enc_len)).bool()
+
+        if self.num_decoder_layers > 0:
+            tgt = self.custom_tgt_module(tgt)
+        if self.attention_type == "RelPosMHAXL":
+            tgt = tgt + self.positional_encoding_decoder(tgt)
+            pos_embs_encoder = None  # self.positional_encoding(src)
+            pos_embs_target = None
+        elif (
+            self.positional_encoding_type == "fixed_abs_sine"
+            or self.attention_type == "hypermixing"
+        ):
+            tgt = tgt + self.positional_encoding(tgt)  # add the encodings here
+            pos_embs_target = None
+            pos_embs_encoder = None
+
+   
+        prediction, self_attns, multihead_attns = self.decoder(
+            tgt,
+            encoder_out,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_key_padding_mask,
+            pos_embs_tgt=pos_embs_target,
+            pos_embs_src=pos_embs_encoder,
+        )
+        return prediction, multihead_attns[-1]
+
+    def encode(
+        self,
+        src,
+        wav_len=None,
+        pad_idx=0,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
+        """
+        Encoder forward pass
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence to the encoder.
+        wav_len : torch.Tensor, optional
+            Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
+        pad_idx : int
+            The index used for padding.
+        dynchunktrain_config : DynChunkTrainConfig
+            Dynamic chunking config.
+
+        Returns
+        -------
+        encoder_out : torch.Tensor
+        """
+        # reshape the src vector to [Batch, Time, Fea] if a 4d vector is given
+        if src.dim() == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        (
+            src_key_padding_mask,
+            _,
+            src_mask,
+            _,
+        ) = make_transformer_src_tgt_masks(
+            src,
+            None,
+            wav_len,
+            pad_idx=pad_idx,
+            causal=self.causal,
+            dynchunktrain_config=dynchunktrain_config,
+        )
+
+        src = self.custom_src_module(src)
+        if self.attention_type == "hypermixing":
+            pos_embs_source = None
+        elif self.attention_type == "RelPosMHAXL":
+            pos_embs_source = self.positional_encoding(src)
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(src)
+            pos_embs_source = None
+
+        encoder_out, _ = self.encoder(
+            src=src,
+            src_mask=src_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            pos_embs=pos_embs_source,
+            dynchunktrain_config=dynchunktrain_config,
+        )
+
+        return encoder_out
+
+    def encode_streaming(self, src, context: TransformerASRStreamingContext):
+        """
+        Streaming encoder forward pass
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence (chunk) to the encoder.
+        context : TransformerASRStreamingContext
+            Mutable reference to the streaming context. This holds the state
+            needed to persist across chunk inferences and can be built using
+            `make_streaming_context`. This will get mutated by this function.
+
+        Returns
+        -------
+        Encoder output for this chunk.
+
+        Example
+        -------
+        >>> import torch
+        >>> from speechbrain.lobes.models.transformer.TransformerASR import TransformerASR
+        >>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+        >>> net = TransformerASR(
+        ...     tgt_vocab=100,
+        ...     input_size=64,
+        ...     d_model=64,
+        ...     nhead=8,
+        ...     num_encoder_layers=1,
+        ...     num_decoder_layers=0,
+        ...     d_ffn=128,
+        ...     attention_type="RelPosMHAXL",
+        ...     positional_encoding=None,
+        ...     encoder_module="conformer",
+        ...     normalize_before=True,
+        ...     causal=False,
+        ... )
+        >>> ctx = net.make_streaming_context(DynChunkTrainConfig(16, 1))
+        >>> src1 = torch.rand([8, 16, 64])
+        >>> src2 = torch.rand([8, 16, 64])
+        >>> out1 = net.encode_streaming(src1, ctx)
+        >>> out1.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 16, 64])
+        >>> out2 = net.encode_streaming(src2, ctx)
+        >>> out2.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 16, 64])
+        >>> combined_out = torch.concat((out1, out2), dim=1)
+        >>> combined_out.shape
+        torch.Size([8, 32, 64])
+        """
+
+        if src.dim() == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        # HACK: our problem here is that the positional_encoding is computed
+        # against the size of our source tensor, but we only know how many left
+        # context frames we're injecting to the encoder within the encoder
+        # context.
+        # so this workaround does just that.
+        #
+        # i'm not sure how this would be best refactored, but an option would be
+        # to let the encoder get the pos embedding itself and have a way to
+        # cache it.
+        #
+        # additionally, positional encoding functions take in a whole source
+        # tensor just to get its attributes (size, device, type) but this is
+        # sort of silly for the embeddings that don't need one.
+        # so we craft a dummy empty (uninitialized) tensor to help...
+        known_left_context = context.encoder_context.layers[0].mha_left_context
+        if known_left_context is None:
+            pos_encoding_dummy = src
+        else:
+            target_shape = list(src.shape)
+            target_shape[-2] += known_left_context.shape[-2]
+            pos_encoding_dummy = torch.empty(size=target_shape).to(src)
+
+        src = self.custom_src_module(src)
+        if self.attention_type == "RelPosMHAXL":
+            pos_embs_source = self.positional_encoding(pos_encoding_dummy)
+
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(pos_encoding_dummy)
+            pos_embs_source = None
+
+        encoder_out, _ = self.encoder.forward_streaming(
+            src=src, pos_embs=pos_embs_source, context=context.encoder_context
+        )
+        return encoder_out
+
+    def make_streaming_context(
+        self, dynchunktrain_config: DynChunkTrainConfig, encoder_kwargs={}
+    ):
+        """Creates a blank streaming context for this transformer and its
+        encoder.
+
+        Arguments
+        ---------
+        dynchunktrain_config : DynChunkTrainConfig
+            Runtime chunkwise attention configuration.
+        encoder_kwargs : dict
+            Parameters to be forward to the encoder's `make_streaming_context`.
+            Metadata required for the encoder could differ depending on the
+            encoder.
+
+        Returns
+        -------
+        TransformerASRStreamingContext
+        """
+        return TransformerASRStreamingContext(
+            dynchunktrain_config=dynchunktrain_config,
+            encoder_context=self.encoder.make_streaming_context(
+                dynchunktrain_config,
+                **encoder_kwargs,
+            ),
+        )
+
+    def _init_params(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                torch.nn.init.xavier_normal_(p)
+
+
+class EncoderWrapper(nn.Module):
+    """This is a wrapper of any ASR transformer encoder. By default, the
+    TransformerASR .forward() function encodes and decodes. With this wrapper
+    the .forward() function becomes .encode() only.
+
+    Important: The TransformerASR class must contain a .encode() function.
+
+    Arguments
+    ---------
+    transformer : sb.lobes.models.TransformerInterface
+        A Transformer instance that contains a .encode() function.
+    *args : tuple
+    **kwargs : dict
+        Arguments to forward to parent class.
+
+    Example
+    -------
+    >>> src = torch.rand([8, 120, 512])
+    >>> tgt = torch.randint(0, 720, [8, 120])
+    >>> net = TransformerASR(
+    ...     720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU
+    ... )
+    >>> encoder = EncoderWrapper(net)
+    >>> enc_out = encoder(src)
+    >>> enc_out.shape
+    torch.Size([8, 120, 512])
+    """
+
+    def __init__(self, transformer, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.transformer = transformer
+        self.make_streaming_context = self.transformer.make_streaming_context
+
+    def forward(self, x, wav_lens=None, pad_idx=0, **kwargs):
+        """Processes the input tensor x and returns an output tensor."""
+        x = self.transformer.encode(x, wav_lens, pad_idx, **kwargs)
+        return x
+
+    def forward_streaming(self, x, context):
+        """Processes the input audio chunk tensor `x`, using and updating the
+        mutable encoder `context`"""
+        x = self.transformer.encode_streaming(x, context)
+        return x
+
+    def make_streaming_context(self, *args, **kwargs):
+        """Initializes a streaming context. Forwards all arguments to the
+        underlying transformer. See :meth:`speechbrain.lobes.models.transformer.TransformerASR.make_streaming_context`.
+        """
+        return self.transformer.make_streaming_context(*args, **kwargs)

modules/mamba/bimamba.py

@@ -0,0 +1,465 @@
+'''
+Copied and modified from 
+https://github.com/hustvl/Vim/blob/main/mamba-1p1p1/mamba_ssm/modules/mamba_simple.py
+'''
+
+# Copyright (c) 2023, Tri Dao, Albert Gu.
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from einops import rearrange, repeat
+
+try:
+    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+except ImportError:
+    causal_conv1d_fn, causal_conv1d_update = None
+
+try:
+    from modules.mamba.selective_scan_interface import selective_scan_fn, mamba_inner_fn, bimamba_inner_fn, mamba_inner_fn_no_out_proj
+except ImportError:
+    selective_scan_fn, mamba_inner_fn, bimamba_inner_fn, mamba_inner_fn_no_out_proj = None, None, None, None, None
+
+try:
+    from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+except ImportError:
+    selective_state_update = None
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+
+class Mamba(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_state=16,
+        d_conv=4,
+        expand=2,
+        dt_rank="auto",
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init="random",
+        dt_scale=1.0,
+        dt_init_floor=1e-4,
+        conv_bias=True,
+        bias=False,
+        use_fast_path=True,  # Fused kernel options
+        layer_idx=None,
+        device=None,
+        dtype=None,
+        bimamba_type="none",
+        if_devide_out=True, # False
+        init_layer_scale=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.expand = expand
+        self.d_inner = int(self.expand * self.d_model)
+        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
+        self.use_fast_path = use_fast_path
+        self.layer_idx = layer_idx
+        self.bimamba_type = bimamba_type
+        self.if_devide_out = if_devide_out
+
+        assert bimamba_type == 'v2'
+
+        self.init_layer_scale = init_layer_scale
+        if init_layer_scale is not None:
+            self.gamma = nn.Parameter(init_layer_scale * torch.ones((d_model)), requires_grad=True)
+
+        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.d_inner,
+            out_channels=self.d_inner,
+            bias=conv_bias,
+            kernel_size=d_conv,
+            groups=self.d_inner,
+            padding=d_conv - 1,
+            **factory_kwargs,
+        )
+
+        self.activation = "silu"
+        self.act = nn.SiLU()
+
+        self.x_proj = nn.Linear(
+            self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
+        )
+        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True, **factory_kwargs)
+
+        # Initialize special dt projection to preserve variance at initialization
+        dt_init_std = self.dt_rank**-0.5 * dt_scale
+        if dt_init == "constant":
+            nn.init.constant_(self.dt_proj.weight, dt_init_std)
+        elif dt_init == "random":
+            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+
+        # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
+        dt = torch.exp(
+            torch.rand(self.d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(inv_dt)
+        # Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
+        self.dt_proj.bias._no_reinit = True
+
+        # S4D real initialization
+        A = repeat(
+            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
+            "n -> d n",
+            d=self.d_inner,
+        ).contiguous()
+        A_log = torch.log(A)  # Keep A_log in fp32
+        self.A_log = nn.Parameter(A_log)
+        self.A_log._no_weight_decay = True
+
+        # D "skip" parameter
+        self.D = nn.Parameter(torch.ones(self.d_inner, device=device))  # Keep in fp32
+        self.D._no_weight_decay = True
+
+        # bidirectional
+        if bimamba_type == "v1":
+            A_b = repeat(
+                torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
+                "n -> d n",
+                d=self.d_inner,
+            ).contiguous()
+            A_b_log = torch.log(A_b)  # Keep A_b_log in fp32
+            self.A_b_log = nn.Parameter(A_b_log)
+            self.A_b_log._no_weight_decay = True
+        elif bimamba_type == "v2":
+            A_b = repeat(
+                torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
+                "n -> d n",
+                d=self.d_inner,
+            ).contiguous()
+            A_b_log = torch.log(A_b)  # Keep A_b_log in fp32
+            self.A_b_log = nn.Parameter(A_b_log)
+            self.A_b_log._no_weight_decay = True 
+
+            self.conv1d_b = nn.Conv1d(
+                in_channels=self.d_inner,
+                out_channels=self.d_inner,
+                bias=conv_bias,
+                kernel_size=d_conv,
+                groups=self.d_inner,
+                padding=d_conv - 1,
+                **factory_kwargs,
+            )
+
+            self.x_proj_b = nn.Linear(
+                self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
+            )
+            self.dt_proj_b = nn.Linear(self.dt_rank, self.d_inner, bias=True, **factory_kwargs)
+
+            self.D_b = nn.Parameter(torch.ones(self.d_inner, device=device))  # Keep in fp32
+            self.D_b._no_weight_decay = True
+
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+
+    def forward(self, hidden_states, inference_params=None):
+        """
+        hidden_states: (B, L, D)
+        Returns: same shape as hidden_states
+        """
+        batch, seqlen, dim = hidden_states.shape
+        conv_state, ssm_state = None, None
+
+        if inference_params is not None:
+            conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
+            if inference_params.seqlen_offset > 0:
+                # The states are updated inplace
+                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
+                return out
+
+        # We do matmul and transpose BLH -> HBL at the same time
+        xz = rearrange(
+            self.in_proj.weight @ rearrange(hidden_states, "b l d -> d (b l)"),
+            "d (b l) -> b d l",
+            l=seqlen,
+        )
+        if self.in_proj.bias is not None:
+            xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), "d -> d 1")
+
+        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
+        # In the backward pass we write dx and dz next to each other to avoid torch.cat
+        if self.use_fast_path and inference_params is None:  # Doesn't support outputting the states
+            if self.bimamba_type == "v1":
+                A_b = -torch.exp(self.A_b_log.float())
+                out = bimamba_inner_fn(
+                    xz,
+                    self.conv1d.weight,
+                    self.conv1d.bias,
+                    self.x_proj.weight,
+                    self.dt_proj.weight,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    A,
+                    A_b,
+                    None,  # input-dependent B
+                    None,  # input-dependent C
+                    self.D.float(),
+                    delta_bias=self.dt_proj.bias.float(),
+                    delta_softplus=True,
+                )    
+            elif self.bimamba_type == "v2":
+                A_b = -torch.exp(self.A_b_log.float())
+                out = mamba_inner_fn_no_out_proj(
+                    xz,
+                    self.conv1d.weight,
+                    self.conv1d.bias,
+                    self.x_proj.weight,
+                    self.dt_proj.weight,
+                    A,
+                    None,  # input-dependent B
+                    None,  # input-dependent C
+                    self.D.float(),
+                    delta_bias=self.dt_proj.bias.float(),
+                    delta_softplus=True,
+                )
+                out_b = mamba_inner_fn_no_out_proj(
+                    xz.flip([-1]),
+                    self.conv1d_b.weight,
+                    self.conv1d_b.bias,
+                    self.x_proj_b.weight,
+                    self.dt_proj_b.weight,
+                    A_b,
+                    None,
+                    None,
+                    self.D_b.float(),
+                    delta_bias=self.dt_proj_b.bias.float(),
+                    delta_softplus=True,
+                )
+                
+                if not self.if_devide_out:
+                    out = F.linear(rearrange(out + out_b.flip([-1]), "b d l -> b l d"), self.out_proj.weight, self.out_proj.bias)
+                else:
+                    out = F.linear(rearrange(0.5*out + 0.5*out_b.flip([-1]), "b d l -> b l d"), self.out_proj.weight, self.out_proj.bias)
+
+            else:
+                out = mamba_inner_fn(
+                    xz,
+                    self.conv1d.weight,
+                    self.conv1d.bias,
+                    self.x_proj.weight,
+                    self.dt_proj.weight,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    A,
+                    None,  # input-dependent B
+                    None,  # input-dependent C
+                    self.D.float(),
+                    delta_bias=self.dt_proj.bias.float(),
+                    delta_softplus=True,
+                )
+        else:
+            x, z = xz.chunk(2, dim=1)
+            # Compute short convolution
+            if conv_state is not None:
+                # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
+                # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
+                conv_state.copy_(F.pad(x, (self.d_conv - x.shape[-1], 0)))  # Update state (B D W)
+            if causal_conv1d_fn is None:
+                x = self.act(self.conv1d(x)[..., :seqlen])
+            else:
+                assert self.activation in ["silu", "swish"]
+                x = causal_conv1d_fn(
+                    x=x,
+                    weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                )
+
+            # We're careful here about the layout, to avoid extra transposes.
+            # We want dt to have d as the slowest moving dimension
+            # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+            x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))  # (bl d)
+            dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
+            dt = self.dt_proj.weight @ dt.t()
+            dt = rearrange(dt, "d (b l) -> b d l", l=seqlen)
+            B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
+            C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
+            assert self.activation in ["silu", "swish"]
+            y = selective_scan_fn(
+                x,
+                dt,
+                A,
+                B,
+                C,
+                self.D.float(),
+                z=z,
+                delta_bias=self.dt_proj.bias.float(),
+                delta_softplus=True,
+                return_last_state=ssm_state is not None,
+            )
+            if ssm_state is not None:
+                y, last_state = y
+                ssm_state.copy_(last_state)
+            y = rearrange(y, "b d l -> b l d")
+            out = self.out_proj(y)
+        if self.init_layer_scale is not None:
+                out = out * self.gamma    
+        return out
+
+    def step(self, hidden_states, conv_state, ssm_state):
+        dtype = hidden_states.dtype
+        assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
+        xz = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
+        x, z = xz.chunk(2, dim=-1)  # (B D)
+
+        # Conv step
+        if causal_conv1d_update is None:
+            conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))  # Update state (B D W)
+            conv_state[:, :, -1] = x
+            x = torch.sum(conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1)  # (B D)
+            if self.conv1d.bias is not None:
+                x = x + self.conv1d.bias
+            x = self.act(x).to(dtype=dtype)
+        else:
+            x = causal_conv1d_update(
+                x,
+                conv_state,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+        x_db = self.x_proj(x)  # (B dt_rank+2*d_state)
+        dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1)
+        # Don't add dt_bias here
+        dt = F.linear(dt, self.dt_proj.weight)  # (B d_inner)
+        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
+
+        # SSM step
+        if selective_state_update is None:
+            # Discretize A and B
+            dt = F.softplus(dt + self.dt_proj.bias.to(dtype=dt.dtype))
+            dA = torch.exp(torch.einsum("bd,dn->bdn", dt, A))
+            dB = torch.einsum("bd,bn->bdn", dt, B)
+            ssm_state.copy_(ssm_state * dA + rearrange(x, "b d -> b d 1") * dB)
+            y = torch.einsum("bdn,bn->bd", ssm_state.to(dtype), C)
+            y = y + self.D.to(dtype) * x
+            y = y * self.act(z)  # (B D)
+        else:
+            y = selective_state_update(
+                ssm_state, x, dt, A, B, C, self.D, z=z, dt_bias=self.dt_proj.bias, dt_softplus=True
+            )
+
+        out = self.out_proj(y)
+        return out.unsqueeze(1), conv_state, ssm_state
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        device = self.out_proj.weight.device
+        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+        conv_state = torch.zeros(
+            batch_size, self.d_model * self.expand, self.d_conv, device=device, dtype=conv_dtype
+        )
+        ssm_dtype = self.dt_proj.weight.dtype if dtype is None else dtype
+        # ssm_dtype = torch.float32
+        ssm_state = torch.zeros(
+            batch_size, self.d_model * self.expand, self.d_state, device=device, dtype=ssm_dtype
+        )
+        return conv_state, ssm_state
+
+    def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+        assert self.layer_idx is not None
+        if self.layer_idx not in inference_params.key_value_memory_dict:
+            batch_shape = (batch_size,)
+            conv_state = torch.zeros(
+                batch_size,
+                self.d_model * self.expand,
+                self.d_conv,
+                device=self.conv1d.weight.device,
+                dtype=self.conv1d.weight.dtype,
+            )
+            ssm_state = torch.zeros(
+                batch_size,
+                self.d_model * self.expand,
+                self.d_state,
+                device=self.dt_proj.weight.device,
+                dtype=self.dt_proj.weight.dtype,
+                # dtype=torch.float32,
+            )
+            inference_params.key_value_memory_dict[self.layer_idx] = (conv_state, ssm_state)
+        else:
+            conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_idx]
+            # TODO: What if batch size changes between generation, and we reuse the same states?
+            if initialize_states:
+                conv_state.zero_()
+                ssm_state.zero_()
+        return conv_state, ssm_state
+
+
+class Block(nn.Module):
+    def __init__(
+        self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False
+    ):
+        """
+        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"
+
+        This Block has a slightly different structure compared to a regular
+        prenorm Transformer block.
+        The standard block is: LN -> MHA/MLP -> Add.
+        [Ref: https://arxiv.org/abs/2002.04745]
+        Here we have: Add -> LN -> Mixer, returning both
+        the hidden_states (output of the mixer) and the residual.
+        This is purely for performance reasons, as we can fuse add and LayerNorm.
+        The residual needs to be provided (except for the very first block).
+        """
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.mixer = mixer_cls(dim)
+        self.norm = norm_cls(dim)
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import fails"
+            assert isinstance(
+                self.norm, (nn.LayerNorm, RMSNorm)
+            ), "Only LayerNorm and RMSNorm are supported for fused_add_norm"
+
+    def forward(
+        self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None
+    ):
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: hidden_states = Mixer(LN(residual))
+        """
+        if not self.fused_add_norm:
+            residual = (hidden_states + residual) if residual is not None else hidden_states
+            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+        else:
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
+            hidden_states, residual = fused_add_norm_fn(
+                hidden_states,
+                self.norm.weight,
+                self.norm.bias,
+                residual=residual,
+                prenorm=True,
+                residual_in_fp32=self.residual_in_fp32,
+                eps=self.norm.eps,
+            )
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        return hidden_states, residual
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)

modules/mamba/mamba_blocks.py

@@ -0,0 +1,252 @@
+'''
+Copied and modified from
+https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py
+'''
+
+import math
+import torch
+import torch.nn as nn
+
+from functools import partial
+
+from mamba_ssm import Mamba
+from modules.mamba.bimamba import Mamba as BiMamba 
+from modules.mamba.bimamba import Block as PreNormBlock
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+
+def create_block(
+    d_model,
+    ssm_cls=None,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=True,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+):
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    factory_kwargs = {"device": device, "dtype": dtype}
+    mixer_cls = partial(ssm_cls, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+    norm_cls = partial(
+        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
+    )
+    block = PreNormBlock(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+class LnMambaAdd(nn.Module):
+
+    def __init__(self, 
+        d_model, 
+        ssm_cls, 
+        ssm_cfg,
+        rms_norm=False,
+        layer_idx=None
+    ):
+        super().__init__()
+        if rms_norm:
+            self.norm = RMSNorm(d_model)
+        else:
+            self.norm = nn.LayerNorm(d_model)
+        self.mamba = ssm_cls(d_model=d_model, **ssm_cfg)
+
+        print(type(self.mamba))
+
+        print('Created LnMambaAdd.')
+
+    def forward(self, x, residual=None, inference_params=None):
+        if residual != None:
+            x = x + residual
+        return self.mamba(self.norm(x)), x
+
+
+class MambaBlocksSequential(nn.Module):
+    """
+    A wrapper for the Mamba block to replicate it
+
+    Arguments
+    ---------
+    n_mamba : int
+        Number of Mamba blocks
+    d_model : int
+        Input dimension to Mamba (bottleneck dimension).
+    d_state : int
+        Mamba state dimension
+    expand: int
+        First linear projection d_model -> d_model * expand
+    d_conv: int
+        Kernel size of Mamba conv
+    norm type : str
+        The type of normalization, in ['gLN', 'cLN'].
+    ---------
+    """
+
+    def __init__(self, 
+        n_mamba: int,
+        bidirectional: bool,
+        d_model: int, # bottleneck dimension (B)
+        d_state: int = 16,
+        expand: int = 2,
+        d_conv: int = 4, # kernel_size of 'Conv' in Mamba
+        dt_rank: str="auto",
+        conv_bias: bool = True,
+        bias: bool = False,
+        fused_add_norm: bool = True,
+        rms_norm: bool = False,
+        norm_epsilon: float = 1e-5,
+        initializer_cfg=None,
+        residual_in_fp32=False,
+        use_simple_block=False
+    ):
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.bidirectional = bidirectional
+
+        # We change the order of residual and layer norm:
+        # Instead of LN -> Attn / MLP -> Add, we do:
+        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
+        # the main branch (output of MLP / Mixer). The model definition is unchanged.
+        # This is for performance reason: we can fuse add + layer_norm.
+        self.fused_add_norm = fused_add_norm
+        if self.fused_add_norm:
+            if layer_norm_fn is None or rms_norm_fn is None:
+                raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")
+
+        self.use_simple_block = use_simple_block
+
+        ssm_cfg = {
+            "d_state": d_state,
+            "expand": expand,
+            "d_conv": d_conv,
+            "dt_rank": dt_rank,
+            "conv_bias": conv_bias,
+            "bias": bias
+        }
+        if bidirectional:
+            ssm_cfg["bimamba_type"] = "v2"
+
+        if use_simple_block:
+            self.layers = nn.Sequential(
+                *[
+                    LnMambaAdd(
+                        d_model=d_model,
+                        ssm_cls=BiMamba if bidirectional else Mamba,
+                        ssm_cfg=ssm_cfg,
+                        rms_norm=rms_norm,
+                        layer_idx=i
+                    )
+                    for i in range(n_mamba)
+                ]
+            )
+        else:
+            self.layers = nn.Sequential(
+                *[
+                    create_block(
+                        d_model=d_model,
+                        ssm_cls=BiMamba if bidirectional else Mamba,
+                        ssm_cfg=ssm_cfg,
+                        norm_epsilon=norm_epsilon,
+                        rms_norm=rms_norm,
+                        residual_in_fp32=residual_in_fp32,
+                        fused_add_norm=fused_add_norm,
+                        layer_idx=i,
+                    )
+                    for i in range(n_mamba)
+                ]
+            )
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            d_model, eps=norm_epsilon
+        )
+
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_mamba,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: block.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+    
+    def forward(self, x, inference_params=None):
+
+        hidden_states = x
+        residual = None
+        for i, layer in enumerate(self.layers):
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+            
+        if not self.fused_add_norm:
+            residual = (hidden_states + residual) if residual is not None else hidden_states
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm=False here since we don't need the residual
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+
+            hidden_states = fused_add_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+
+        return hidden_states
+        

modules/mamba/selective_scan_interface.py

@@ -0,0 +1,714 @@
+'''
+Copied from
+https://github.com/hustvl/Vim/blob/main/mamba-1p1p1/mamba_ssm/ops/selective_scan_interface.py
+'''
+
+# Copyright (c) 2023, Tri Dao, Albert Gu.
+
+import torch
+import torch.nn.functional as F
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+from einops import rearrange, repeat
+
+from causal_conv1d import causal_conv1d_fn
+import causal_conv1d_cuda
+import selective_scan_cuda
+
+
+class SelectiveScanFn(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
+                return_last_state=False):
+        if u.stride(-1) != 1:
+            u = u.contiguous()
+        if delta.stride(-1) != 1:
+            delta = delta.contiguous()
+        if D is not None:
+            D = D.contiguous()
+        if B.stride(-1) != 1:
+            B = B.contiguous()
+        if C.stride(-1) != 1:
+            C = C.contiguous()
+        if z is not None and z.stride(-1) != 1:
+            z = z.contiguous()
+        if B.dim() == 3:
+            B = rearrange(B, "b dstate l -> b 1 dstate l")
+            ctx.squeeze_B = True
+        if C.dim() == 3:
+            C = rearrange(C, "b dstate l -> b 1 dstate l")
+            ctx.squeeze_C = True
+        out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, z, delta_bias, delta_softplus)
+        ctx.delta_softplus = delta_softplus
+        ctx.has_z = z is not None
+        last_state = x[:, :, -1, 1::2]  # (batch, dim, dstate)
+        if not ctx.has_z:
+            ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
+            return out if not return_last_state else (out, last_state)
+        else:
+            ctx.save_for_backward(u, delta, A, B, C, D, z, delta_bias, x, out)
+            out_z = rest[0]
+            return out_z if not return_last_state else (out_z, last_state)
+
+    @staticmethod
+    def backward(ctx, dout, *args):
+        if not ctx.has_z:
+            u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
+            z = None
+            out = None
+        else:
+            u, delta, A, B, C, D, z, delta_bias, x, out = ctx.saved_tensors
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the
+        # backward of selective_scan_cuda with the backward of chunk).
+        # Here we just pass in None and dz will be allocated in the C++ code.
+        du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd(
+            u, delta, A, B, C, D, z, delta_bias, dout, x, out, None, ctx.delta_softplus,
+            False  # option to recompute out_z, not used here
+        )
+        dz = rest[0] if ctx.has_z else None
+        dB = dB.squeeze(1) if getattr(ctx, "squeeze_B", False) else dB
+        dC = dC.squeeze(1) if getattr(ctx, "squeeze_C", False) else dC
+        return (du, ddelta, dA, dB, dC,
+                dD if D is not None else None,
+                dz,
+                ddelta_bias if delta_bias is not None else None,
+                None,
+                None)
+
+
+def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
+                     return_last_state=False):
+    """if return_last_state is True, returns (out, last_state)
+    last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
+    not considered in the backward pass.
+    """
+    return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
+
+
+def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
+                      return_last_state=False):
+    """
+    u: r(B D L)
+    delta: r(B D L)
+    A: c(D N) or r(D N)
+    B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
+    C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
+    D: r(D)
+    z: r(B D L)
+    delta_bias: r(D), fp32
+
+    out: r(B D L)
+    last_state (optional): r(B D dstate) or c(B D dstate)
+    """
+    dtype_in = u.dtype
+    u = u.float()
+    delta = delta.float()
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None].float()
+    if delta_softplus:
+        delta = F.softplus(delta)
+    batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
+    is_variable_B = B.dim() >= 3
+    is_variable_C = C.dim() >= 3
+    if A.is_complex():
+        if is_variable_B:
+            B = torch.view_as_complex(rearrange(B.float(), "... (L two) -> ... L two", two=2))
+        if is_variable_C:
+            C = torch.view_as_complex(rearrange(C.float(), "... (L two) -> ... L two", two=2))
+    else:
+        B = B.float()
+        C = C.float()
+    x = A.new_zeros((batch, dim, dstate))
+    ys = []
+    deltaA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))
+    if not is_variable_B:
+        deltaB_u = torch.einsum('bdl,dn,bdl->bdln', delta, B, u)
+    else:
+        if B.dim() == 3:
+            deltaB_u = torch.einsum('bdl,bnl,bdl->bdln', delta, B, u)
+        else:
+            B = repeat(B, "B G N L -> B (G H) N L", H=dim // B.shape[1])
+            deltaB_u = torch.einsum('bdl,bdnl,bdl->bdln', delta, B, u)
+    if is_variable_C and C.dim() == 4:
+        C = repeat(C, "B G N L -> B (G H) N L", H=dim // C.shape[1])
+    last_state = None
+    for i in range(u.shape[2]):
+        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
+        if not is_variable_C:
+            y = torch.einsum('bdn,dn->bd', x, C)
+        else:
+            if C.dim() == 3:
+                y = torch.einsum('bdn,bn->bd', x, C[:, :, i])
+            else:
+                y = torch.einsum('bdn,bdn->bd', x, C[:, :, :, i])
+        if i == u.shape[2] - 1:
+            last_state = x
+        if y.is_complex():
+            y = y.real * 2
+        ys.append(y)
+    y = torch.stack(ys, dim=2) # (batch dim L)
+    out = y if D is None else y + u * rearrange(D, "d -> d 1")
+    if z is not None:
+        out = out * F.silu(z)
+    out = out.to(dtype=dtype_in)
+    return out if not return_last_state else (out, last_state)
+
+
+class MambaInnerFnNoOutProj(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+                C_proj_bias=None, delta_softplus=True, checkpoint_lvl=1):
+        """
+             xz: (batch, dim, seqlen)
+        """
+        assert checkpoint_lvl in [0, 1]
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        if torch.is_autocast_enabled():
+            x_proj_weight = x_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            delta_proj_weight = delta_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+        if xz.stride(-1) != 1:
+            xz = xz.contiguous()
+        conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")
+        x, z = xz.chunk(2, dim=1)
+        conv1d_bias = conv1d_bias.contiguous() if conv1d_bias is not None else None
+        conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias,None, True)
+        # We're being very careful here about the layout, to avoid extra transposes.
+        # We want delta to have d as the slowest moving dimension
+        # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+        x_dbl = F.linear(rearrange(conv1d_out, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
+        delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l = L)
+        ctx.is_variable_B = B is None
+        ctx.is_variable_C = C is None
+        ctx.B_proj_bias_is_None = B_proj_bias is None
+        ctx.C_proj_bias_is_None = C_proj_bias is None
+        if B is None:  # variable B
+            B = x_dbl[:, delta_rank:delta_rank + d_state]  # (bl dstate)
+            if B_proj_bias is not None:
+                B = B + B_proj_bias.to(dtype=B.dtype)
+            if not A.is_complex():
+                # B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous()
+                B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                B = rearrange(B, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if B.stride(-1) != 1:
+                B = B.contiguous()
+        if C is None:  # variable C
+            C = x_dbl[:, -d_state:]  # (bl dstate)
+            if C_proj_bias is not None:
+                C = C + C_proj_bias.to(dtype=C.dtype)
+            if not A.is_complex():
+                # C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous()
+                C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                C = rearrange(C, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if C.stride(-1) != 1:
+                C = C.contiguous()
+        if D is not None:
+            D = D.contiguous()
+        out, scan_intermediates, out_z = selective_scan_cuda.fwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
+        )
+        ctx.delta_softplus = delta_softplus
+        ctx.checkpoint_lvl = checkpoint_lvl
+        if checkpoint_lvl >= 1:  # Will recompute conv1d_out and delta in the backward pass
+            conv1d_out, delta = None, None
+        ctx.save_for_backward(xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight,
+                              delta_proj_weight, conv1d_out, delta,
+                              A, B, C, D, delta_bias, scan_intermediates, out)
+        # return rearrange(out_z, "b d l -> b l d")
+        return out_z
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, dout):
+        # dout: (batch, seqlen, dim)
+        (xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight, delta_proj_weight, 
+         conv1d_out, delta, A, B, C, D, delta_bias, scan_intermediates, out) = ctx.saved_tensors
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        x, z = xz.chunk(2, dim=1)
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        if ctx.checkpoint_lvl == 1:
+            conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias,None, True)
+            delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(),
+                              "d (b l) -> b d l", l = L)
+        # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the
+        # backward of selective_scan_cuda with the backward of chunk).
+        dxz = torch.empty_like(xz)  # (batch, dim, seqlen)
+        dx, dz = dxz.chunk(2, dim=1)
+        # dout_y = rearrange(dout, "b l d -> b d l") # because no arrange at end of forward, so dout shape is b d l
+        dconv1d_out, ddelta, dA, dB, dC, dD, ddelta_bias, dz, out_z = selective_scan_cuda.bwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, dout, scan_intermediates, out, dz,
+            ctx.delta_softplus,
+            True  # option to recompute out_z
+        )
+        dD = dD if D is not None else None
+        dx_dbl = torch.empty_like(x_dbl)
+        dB_proj_bias = None
+        if ctx.is_variable_B:
+            if not A.is_complex():
+                dB = rearrange(dB, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dB = rearrange(dB, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None
+            dx_dbl[:, delta_rank:delta_rank + d_state] = dB  # (bl d)
+            dB = None
+        dC_proj_bias = None
+        if ctx.is_variable_C:
+            if not A.is_complex():
+                dC = rearrange(dC, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dC = rearrange(dC, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dC_proj_bias = dC.sum(0) if not ctx.C_proj_bias_is_None else None
+            dx_dbl[:, -d_state:] = dC  # (bl d)
+            dC = None
+        ddelta = rearrange(ddelta, "b d l -> d (b l)")
+        ddelta_proj_weight = torch.einsum("dB,Br->dr", ddelta, x_dbl[:, :delta_rank])
+        dx_dbl[:, :delta_rank] = torch.einsum("dB,dr->Br", ddelta, delta_proj_weight)
+        dconv1d_out = rearrange(dconv1d_out, "b d l -> d (b l)")
+        dx_proj_weight = torch.einsum("Br,Bd->rd", dx_dbl, rearrange(conv1d_out, "b d l -> (b l) d"))
+        dconv1d_out = torch.addmm(dconv1d_out, x_proj_weight.t(), dx_dbl.t(), out=dconv1d_out)
+        dconv1d_out = rearrange(dconv1d_out, "d (b l) -> b d l", b=x.shape[0], l=x.shape[-1])
+        # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the
+        # backward of conv1d with the backward of chunk).
+        dx, dconv1d_weight, dconv1d_bias = causal_conv1d_cuda.causal_conv1d_bwd(
+            x, conv1d_weight, conv1d_bias, dconv1d_out, None, dx, True
+        )
+        dconv1d_bias = dconv1d_bias if conv1d_bias is not None else None
+        dconv1d_weight = rearrange(dconv1d_weight, "d w -> d 1 w")
+        return (dxz, dconv1d_weight, dconv1d_bias, dx_proj_weight, ddelta_proj_weight,
+                dA, dB, dC, dD,
+                ddelta_bias if delta_bias is not None else None,
+                dB_proj_bias, dC_proj_bias, None)
+    
+
+class MambaInnerFn(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                out_proj_weight, out_proj_bias,
+                A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+                C_proj_bias=None, delta_softplus=True, checkpoint_lvl=1):
+        """
+             xz: (batch, dim, seqlen)
+        """
+        assert checkpoint_lvl in [0, 1]
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        if torch.is_autocast_enabled():
+            x_proj_weight = x_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            delta_proj_weight = delta_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            out_proj_weight = out_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            out_proj_bias = (out_proj_bias.to(dtype=torch.get_autocast_gpu_dtype())
+                             if out_proj_bias is not None else None)
+        if xz.stride(-1) != 1:
+            xz = xz.contiguous()
+        conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")
+        x, z = xz.chunk(2, dim=1)
+        conv1d_bias = conv1d_bias.contiguous() if conv1d_bias is not None else None
+        conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias, None, True)
+        # We're being very careful here about the layout, to avoid extra transposes.
+        # We want delta to have d as the slowest moving dimension
+        # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+        x_dbl = F.linear(rearrange(conv1d_out, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
+        delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l = L)
+        ctx.is_variable_B = B is None
+        ctx.is_variable_C = C is None
+        ctx.B_proj_bias_is_None = B_proj_bias is None
+        ctx.C_proj_bias_is_None = C_proj_bias is None
+        if B is None:  # variable B
+            B = x_dbl[:, delta_rank:delta_rank + d_state]  # (bl dstate)
+            if B_proj_bias is not None:
+                B = B + B_proj_bias.to(dtype=B.dtype)
+            if not A.is_complex():
+                # B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous()
+                B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                B = rearrange(B, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if B.stride(-1) != 1:
+                B = B.contiguous()
+        if C is None:  # variable C
+            C = x_dbl[:, -d_state:]  # (bl dstate)
+            if C_proj_bias is not None:
+                C = C + C_proj_bias.to(dtype=C.dtype)
+            if not A.is_complex():
+                # C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous()
+                C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                C = rearrange(C, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if C.stride(-1) != 1:
+                C = C.contiguous()
+        if D is not None:
+            D = D.contiguous()
+        out, scan_intermediates, out_z = selective_scan_cuda.fwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
+        )
+        ctx.delta_softplus = delta_softplus
+        ctx.out_proj_bias_is_None = out_proj_bias is None
+        ctx.checkpoint_lvl = checkpoint_lvl
+        if checkpoint_lvl >= 1:  # Will recompute conv1d_out and delta in the backward pass
+            conv1d_out, delta = None, None
+        ctx.save_for_backward(xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight,
+                              delta_proj_weight, out_proj_weight, conv1d_out, delta,
+                              A, B, C, D, delta_bias, scan_intermediates, out)
+        return F.linear(rearrange(out_z, "b d l -> b l d"), out_proj_weight, out_proj_bias)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, dout):
+        # dout: (batch, seqlen, dim)
+        (xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight, delta_proj_weight, out_proj_weight,
+         conv1d_out, delta, A, B, C, D, delta_bias, scan_intermediates, out) = ctx.saved_tensors
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        x, z = xz.chunk(2, dim=1)
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        if ctx.checkpoint_lvl == 1:
+            conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias, None, True)
+            delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(),
+                              "d (b l) -> b d l", l = L)
+        # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the
+        # backward of selective_scan_cuda with the backward of chunk).
+        dxz = torch.empty_like(xz)  # (batch, dim, seqlen)
+        dx, dz = dxz.chunk(2, dim=1)
+        dout = rearrange(dout, "b l e -> e (b l)")
+        dout_y = rearrange(out_proj_weight.t() @ dout, "d (b l) -> b d l", l=L)
+        dconv1d_out, ddelta, dA, dB, dC, dD, ddelta_bias, dz, out_z = selective_scan_cuda.bwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, dout_y, scan_intermediates, out, dz,
+            ctx.delta_softplus,
+            True  # option to recompute out_z
+        )
+        dout_proj_weight = torch.einsum("eB,dB->ed", dout, rearrange(out_z, "b d l -> d (b l)"))
+        dout_proj_bias = dout.sum(dim=(0, 1)) if not ctx.out_proj_bias_is_None else None
+        dD = dD if D is not None else None
+        dx_dbl = torch.empty_like(x_dbl)
+        dB_proj_bias = None
+        if ctx.is_variable_B:
+            if not A.is_complex():
+                dB = rearrange(dB, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dB = rearrange(dB, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None
+            dx_dbl[:, delta_rank:delta_rank + d_state] = dB  # (bl d)
+            dB = None
+        dC_proj_bias = None
+        if ctx.is_variable_C:
+            if not A.is_complex():
+                dC = rearrange(dC, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dC = rearrange(dC, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dC_proj_bias = dC.sum(0) if not ctx.C_proj_bias_is_None else None
+            dx_dbl[:, -d_state:] = dC  # (bl d)
+            dC = None
+        ddelta = rearrange(ddelta, "b d l -> d (b l)")
+        ddelta_proj_weight = torch.einsum("dB,Br->dr", ddelta, x_dbl[:, :delta_rank])
+        dx_dbl[:, :delta_rank] = torch.einsum("dB,dr->Br", ddelta, delta_proj_weight)
+        dconv1d_out = rearrange(dconv1d_out, "b d l -> d (b l)")
+        dx_proj_weight = torch.einsum("Br,Bd->rd", dx_dbl, rearrange(conv1d_out, "b d l -> (b l) d"))
+        dconv1d_out = torch.addmm(dconv1d_out, x_proj_weight.t(), dx_dbl.t(), out=dconv1d_out)
+        dconv1d_out = rearrange(dconv1d_out, "d (b l) -> b d l", b=x.shape[0], l=x.shape[-1])
+        # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the
+        # backward of conv1d with the backward of chunk).
+        dx, dconv1d_weight, dconv1d_bias = causal_conv1d_cuda.causal_conv1d_bwd(
+            x, conv1d_weight, conv1d_bias, dconv1d_out, None, dx, True
+        )
+        dconv1d_bias = dconv1d_bias if conv1d_bias is not None else None
+        dconv1d_weight = rearrange(dconv1d_weight, "d w -> d 1 w")
+        return (dxz, dconv1d_weight, dconv1d_bias, dx_proj_weight, ddelta_proj_weight,
+                dout_proj_weight, dout_proj_bias,
+                dA, dB, dC, dD,
+                ddelta_bias if delta_bias is not None else None,
+                dB_proj_bias, dC_proj_bias, None)
+
+
+class BiMambaInnerFn(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                out_proj_weight, out_proj_bias,
+                A, A_b, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+                C_proj_bias=None, delta_softplus=True, checkpoint_lvl=1):
+        """
+             xz: (batch, dim, seqlen)
+        """
+        assert checkpoint_lvl in [0, 1]
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        if torch.is_autocast_enabled():
+            x_proj_weight = x_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            delta_proj_weight = delta_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            out_proj_weight = out_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
+            out_proj_bias = (out_proj_bias.to(dtype=torch.get_autocast_gpu_dtype())
+                             if out_proj_bias is not None else None)
+        if xz.stride(-1) != 1:
+            xz = xz.contiguous()
+        conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")
+        x, z = xz.chunk(2, dim=1)
+        conv1d_bias = conv1d_bias.contiguous() if conv1d_bias is not None else None
+        conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias,None, True)
+        # We're being very careful here about the layout, to avoid extra transposes.
+        # We want delta to have d as the slowest moving dimension
+        # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+        x_dbl = F.linear(rearrange(conv1d_out, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
+        delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l = L)
+        ctx.is_variable_B = B is None
+        ctx.is_variable_C = C is None
+        ctx.B_proj_bias_is_None = B_proj_bias is None
+        ctx.C_proj_bias_is_None = C_proj_bias is None
+        if B is None:  # variable B
+            B = x_dbl[:, delta_rank:delta_rank + d_state]  # (bl dstate)
+            if B_proj_bias is not None:
+                B = B + B_proj_bias.to(dtype=B.dtype)
+            if not A.is_complex():
+                # B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous()
+                B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                B = rearrange(B, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if B.stride(-1) != 1:
+                B = B.contiguous()
+        if C is None:  # variable C
+            C = x_dbl[:, -d_state:]  # (bl dstate)
+            if C_proj_bias is not None:
+                C = C + C_proj_bias.to(dtype=C.dtype)
+            if not A.is_complex():
+                # C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous()
+                C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
+            else:
+                C = rearrange(C, "(b l) (dstate two) -> b 1 dstate (l two)", l=L, two=2).contiguous()
+        else:
+            if C.stride(-1) != 1:
+                C = C.contiguous()
+        if D is not None:
+            D = D.contiguous()
+        out_f, scan_intermediates_f, out_z_f = selective_scan_cuda.fwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
+        )
+        assert not A_b.is_complex(), "A should not be complex!!"
+        out_b, scan_intermediates_b, out_z_b = selective_scan_cuda.fwd(
+            conv1d_out.flip([-1]), delta.flip([-1]), A_b, B.flip([-1]), C.flip([-1]), D, z.flip([-1]), delta_bias, delta_softplus,
+        )
+
+        out_z = out_z_f + out_z_b.flip([-1])
+
+        ctx.delta_softplus = delta_softplus
+        ctx.out_proj_bias_is_None = out_proj_bias is None
+        ctx.checkpoint_lvl = checkpoint_lvl
+        if checkpoint_lvl >= 1:  # Will recompute conv1d_out and delta in the backward pass
+            conv1d_out, delta = None, None
+        ctx.save_for_backward(xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight,
+                              delta_proj_weight, out_proj_weight, conv1d_out, delta,
+                              A, A_b, B, C, D, delta_bias, scan_intermediates_f, scan_intermediates_b, out_f, out_b)
+        return F.linear(rearrange(out_z, "b d l -> b l d"), out_proj_weight, out_proj_bias)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, dout):
+        # dout: (batch, seqlen, dim)
+        (xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight, delta_proj_weight, out_proj_weight,
+         conv1d_out, delta, A, A_b, B, C, D, delta_bias, scan_intermediates_f, scan_intermediates_b, out_f, out_b) = ctx.saved_tensors
+        L = xz.shape[-1]
+        delta_rank = delta_proj_weight.shape[1]
+        d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+        x, z = xz.chunk(2, dim=1)
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        if ctx.checkpoint_lvl == 1:
+            conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias, None, True)
+            delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(),
+                              "d (b l) -> b d l", l = L)
+        # The kernel supports passing in a pre-allocated dz (e.g., in case we want to fuse the
+        # backward of selective_scan_cuda with the backward of chunk).
+        dxz = torch.empty_like(xz)  # (batch, dim, seqlen)
+        dx, dz = dxz.chunk(2, dim=1)
+        dout = rearrange(dout, "b l e -> e (b l)")
+        dout_y = rearrange(out_proj_weight.t() @ dout, "d (b l) -> b d l", l=L)
+        dconv1d_out, ddelta, dA, dB, dC, dD, ddelta_bias, dz, out_z_f = selective_scan_cuda.bwd(
+            conv1d_out, delta, A, B, C, D, z, delta_bias, dout_y, scan_intermediates_f, out_f, dz,
+            ctx.delta_softplus,
+            True  # option to recompute out_z
+        )
+        # flip one
+        dz_b = torch.empty_like(dz)
+        dconv1d_out_f_b, ddelta_f_b, dA_b, dB_f_b, dC_f_b, dD_b, ddelta_bias_b, dz_b, out_z_b = selective_scan_cuda.bwd(
+            conv1d_out.flip([-1]), delta.flip([-1]), A_b, B.flip([-1]), C.flip([-1]), D, z.flip([-1]), delta_bias, dout_y.flip([-1]), scan_intermediates_b, out_b, dz_b,
+            ctx.delta_softplus,
+            True  # option to recompute out_z
+        )
+
+        dconv1d_out = dconv1d_out + dconv1d_out_f_b.flip([-1])
+        ddelta = ddelta + ddelta_f_b.flip([-1])
+        dB = dB + dB_f_b.flip([-1])
+        dC = dC + dC_f_b.flip([-1])
+        dD = dD + dD_b
+        ddelta_bias = ddelta_bias + ddelta_bias_b
+        dz = dz + dz_b.flip([-1])
+        out_z = out_z_f + out_z_b.flip([-1])
+        
+        dout_proj_weight = torch.einsum("eB,dB->ed", dout, rearrange(out_z, "b d l -> d (b l)"))
+        dout_proj_bias = dout.sum(dim=(0, 1)) if not ctx.out_proj_bias_is_None else None
+        dD = dD if D is not None else None
+        dx_dbl = torch.empty_like(x_dbl)
+        dB_proj_bias = None
+        if ctx.is_variable_B:
+            if not A.is_complex():
+                dB = rearrange(dB, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dB = rearrange(dB, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None
+            dx_dbl[:, delta_rank:delta_rank + d_state] = dB  # (bl d)
+            dB = None
+        dC_proj_bias = None
+        if ctx.is_variable_C:
+            if not A.is_complex():
+                dC = rearrange(dC, "b 1 dstate l -> (b l) dstate").contiguous()
+            else:
+                dC = rearrange(dC, "b 1 dstate (l two) -> (b l) (dstate two)", two=2).contiguous()
+            dC_proj_bias = dC.sum(0) if not ctx.C_proj_bias_is_None else None
+            dx_dbl[:, -d_state:] = dC  # (bl d)
+            dC = None
+        ddelta = rearrange(ddelta, "b d l -> d (b l)")
+        ddelta_proj_weight = torch.einsum("dB,Br->dr", ddelta, x_dbl[:, :delta_rank])
+        dx_dbl[:, :delta_rank] = torch.einsum("dB,dr->Br", ddelta, delta_proj_weight)
+        dconv1d_out = rearrange(dconv1d_out, "b d l -> d (b l)")
+        dx_proj_weight = torch.einsum("Br,Bd->rd", dx_dbl, rearrange(conv1d_out, "b d l -> (b l) d"))
+        dconv1d_out = torch.addmm(dconv1d_out, x_proj_weight.t(), dx_dbl.t(), out=dconv1d_out)
+        dconv1d_out = rearrange(dconv1d_out, "d (b l) -> b d l", b=x.shape[0], l=x.shape[-1])
+        # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the
+        # backward of conv1d with the backward of chunk).
+        dx, dconv1d_weight, dconv1d_bias = causal_conv1d_cuda.causal_conv1d_bwd(
+            x, conv1d_weight, conv1d_bias, dconv1d_out, None, dx, True
+        )
+        dconv1d_bias = dconv1d_bias if conv1d_bias is not None else None
+        dconv1d_weight = rearrange(dconv1d_weight, "d w -> d 1 w")
+        return (dxz, dconv1d_weight, dconv1d_bias, dx_proj_weight, ddelta_proj_weight,
+                dout_proj_weight, dout_proj_bias,
+                dA, dA_b, dB, dC, dD,
+                ddelta_bias if delta_bias is not None else None,
+                dB_proj_bias, dC_proj_bias, None)
+    
+
+def mamba_inner_fn(
+    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+    out_proj_weight, out_proj_bias,
+    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+    C_proj_bias=None, delta_softplus=True
+):
+    return MambaInnerFn.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                              out_proj_weight, out_proj_bias,
+                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)
+
+def bimamba_inner_fn(
+    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+    out_proj_weight, out_proj_bias,
+    A, A_b, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+    C_proj_bias=None, delta_softplus=True
+):
+    return BiMambaInnerFn.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                              out_proj_weight, out_proj_bias,
+                              A, A_b, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)
+
+
+def mamba_inner_fn_no_out_proj(
+    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+    C_proj_bias=None, delta_softplus=True
+):
+    return MambaInnerFnNoOutProj.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)
+
+
+def mamba_inner_ref(
+    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+    out_proj_weight, out_proj_bias,
+    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+    C_proj_bias=None, delta_softplus=True
+):
+    L = xz.shape[-1]
+    delta_rank = delta_proj_weight.shape[1]
+    d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+    x, z = xz.chunk(2, dim=1)
+    x = causal_conv1d_fn(x, rearrange(conv1d_weight, "d 1 w -> d w"), conv1d_bias, "silu")
+    # We're being very careful here about the layout, to avoid extra transposes.
+    # We want delta to have d as the slowest moving dimension
+    # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+    x_dbl = F.linear(rearrange(x, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
+    delta = delta_proj_weight @ x_dbl[:, :delta_rank].t()
+    delta = rearrange(delta, "d (b l) -> b d l", l=L)
+    if B is None:  # variable B
+        B = x_dbl[:, delta_rank:delta_rank + d_state]  # (bl d)
+        if B_proj_bias is not None:
+            B = B + B_proj_bias.to(dtype=B.dtype)
+        if not A.is_complex():
+            B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous()
+        else:
+            B = rearrange(B, "(b l) (dstate two) -> b dstate (l two)", l=L, two=2).contiguous()
+    if C is None:  # variable B
+        C = x_dbl[:, -d_state:]  # (bl d)
+        if C_proj_bias is not None:
+            C = C + C_proj_bias.to(dtype=C.dtype)
+        if not A.is_complex():
+            C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous()
+        else:
+            C = rearrange(C, "(b l) (dstate two) -> b dstate (l two)", l=L, two=2).contiguous()
+    y = selective_scan_fn(x, delta, A, B, C, D, z=z, delta_bias=delta_bias, delta_softplus=True)
+    return F.linear(rearrange(y, "b d l -> b l d"), out_proj_weight, out_proj_bias)
+
+
+def bimamba_inner_ref(
+    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
+    out_proj_weight, out_proj_bias,
+    A, A_b, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
+    C_proj_bias=None, delta_softplus=True
+):
+    L = xz.shape[-1]
+    delta_rank = delta_proj_weight.shape[1]
+    d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
+    x, z = xz.chunk(2, dim=1)
+    x = causal_conv1d_fn(x, rearrange(conv1d_weight, "d 1 w -> d w"), conv1d_bias, "silu")
+    # We're being very careful here about the layout, to avoid extra transposes.
+    # We want delta to have d as the slowest moving dimension
+    # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
+    x_dbl = F.linear(rearrange(x, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
+    delta = delta_proj_weight @ x_dbl[:, :delta_rank].t()
+    delta = rearrange(delta, "d (b l) -> b d l", l=L)
+    if B is None:  # variable B
+        B = x_dbl[:, delta_rank:delta_rank + d_state]  # (bl d)
+        if B_proj_bias is not None:
+            B = B + B_proj_bias.to(dtype=B.dtype)
+        if not A.is_complex():
+            B = rearrange(B, "(b l) dstate -> b dstate l", l=L).contiguous()
+        else:
+            B = rearrange(B, "(b l) (dstate two) -> b dstate (l two)", l=L, two=2).contiguous()
+    if C is None:  # variable B
+        C = x_dbl[:, -d_state:]  # (bl d)
+        if C_proj_bias is not None:
+            C = C + C_proj_bias.to(dtype=C.dtype)
+        if not A.is_complex():
+            C = rearrange(C, "(b l) dstate -> b dstate l", l=L).contiguous()
+        else:
+            C = rearrange(C, "(b l) (dstate two) -> b dstate (l two)", l=L, two=2).contiguous()
+    y = selective_scan_fn(x, delta, A, B, C, D, z=z, delta_bias=delta_bias, delta_softplus=True)
+    y_b = selective_scan_fn(x.flip([-1]), delta.flip([-1]), A_b, B.flip([-1]), C.flip([-1]), D, z.flip([-1]), delta_bias, delta_softplus=True)
+    y = y + y_b.flip([-1])
+    return F.linear(rearrange(y, "b d l -> b l d"), out_proj_weight, out_proj_bias)

requirements.txt

@@ -0,0 +1,38 @@
+# Tested with Python 3.10 and CUDA 11.8.
+# The CUDA extensions at the end of this file must be built after PyTorch and
+# the build tools are installed; the model card documents the install order.
+
+--extra-index-url https://download.pytorch.org/whl/cu118
+
+# PyTorch CUDA runtime
+torch==2.0.1+cu118
+torchaudio==2.0.2+cu118
+triton==2.0.0
+
+# Extension build requirements
+setuptools==68.2.2
+wheel==0.41.3
+Cython==0.29.37
+ninja==1.11.1.1
+packaging==24.0
+buildtools==1.0.6
+
+# ConMamba and SpeechBrain runtime
+numpy==1.26.4
+scipy==1.15.3
+soundfile==0.14.0
+einops==0.8.1
+sentencepiece==0.2.0
+speechbrain==1.0.0
+HyperPyYAML==1.2.2
+ruamel.yaml==0.19.1
+huggingface-hub==0.20.3
+joblib==1.5.3
+tqdm==4.70.0
+transformers==4.35.2
+tokenizers==0.15.2
+safetensors==0.8.0
+
+# CUDA extensions used by the bundled Mamba implementation
+causal-conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@89e4aa413ea99a079a39364b0a58f54239b4306c
+mamba-ssm @ git+https://github.com/state-spaces/mamba.git@009bec5ee37f586844a3fc89c040a9c1a9d8badf

run_inference.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Transcribe an audio file with the local ConMamba-small-ca checkpoint."""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torchaudio
+from hyperpyyaml import load_hyperpyyaml
+
+
+MODEL_ROOT = Path(__file__).resolve().parent
+CONFIG_PATH = MODEL_ROOT / "inference.yaml"
+
+
+class ConMambaASR:
+    """Standalone inference wrapper for ConMamba-small-ca."""
+
+    def __init__(self, device: str = "cuda"):
+        self.device = torch.device(device)
+        if self.device.type == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError(
+                "ConMamba-small-ca requires a CUDA-capable GPU for the "
+                "default device, but CUDA is not available."
+            )
+
+        self._validate_repository()
+        self._make_local_modules_importable()
+        hparams = self._load_hparams()
+
+        self.model = hparams["model"]
+        self.compute_features = hparams["compute_features"]
+        self.normalize = hparams["modules"]["normalize"]
+        self.cnn = hparams["modules"]["CNN"]
+        self.transformer = hparams["modules"]["Transformer"]
+        self.search = hparams["beam_search"]
+        self.tokenizer = hparams["tokenizer"]
+        self.bos_index = int(hparams["bos_index"])
+        self.pad_index = int(hparams["pad_index"])
+        self.sample_rate = int(hparams["sample_rate"])
+
+        self._load_artifacts(hparams)
+
+    @staticmethod
+    def _make_local_modules_importable() -> None:
+        model_root = str(MODEL_ROOT)
+        if model_root not in sys.path:
+            sys.path.insert(0, model_root)
+
+    @staticmethod
+    def _validate_repository() -> None:
+        required_paths = [
+            CONFIG_PATH,
+            MODEL_ROOT / "model.ckpt",
+            MODEL_ROOT / "tokenizer.ckpt",
+            MODEL_ROOT / "normalizer.ckpt",
+            MODEL_ROOT / "modules" / "TransformerASR.py",
+        ]
+        missing = [str(path) for path in required_paths if not path.is_file()]
+        if missing:
+            raise FileNotFoundError(
+                "The ConMamba-small-ca repository is incomplete; missing: "
+                + ", ".join(missing)
+            )
+
+    @staticmethod
+    def _load_hparams():
+        # HyperPyYAML 1.2.2 predates ruamel.yaml 0.19, which expects custom
+        # loaders to expose this optional recursion-limit attribute.
+        import ruamel.yaml
+
+        if not hasattr(ruamel.yaml.Loader, "max_depth"):
+            ruamel.yaml.Loader.max_depth = 0
+
+        overrides = {"model_dir": str(MODEL_ROOT)}
+        with CONFIG_PATH.open("r", encoding="utf-8") as config_file:
+            return load_hyperpyyaml(config_file, overrides)
+
+    def _load_artifacts(self, hparams) -> None:
+        model_path = Path(hparams["model_checkpoint"])
+        normalizer_path = Path(hparams["normalizer_checkpoint"])
+        tokenizer_path = Path(hparams["tokenizer_checkpoint"])
+
+        model_state = torch.load(model_path, map_location="cpu")
+        self.model.load_state_dict(model_state, strict=True)
+        self.normalize._load(normalizer_path)
+
+        if not self.tokenizer.load(str(tokenizer_path)):
+            raise RuntimeError(f"Failed to load tokenizer: {tokenizer_path}")
+
+        self.model.to(self.device).eval()
+        self.normalize.eval()
+        if hasattr(self.compute_features, "to"):
+            self.compute_features.to(self.device)
+        if hasattr(self.compute_features, "eval"):
+            self.compute_features.eval()
+
+    def _load_audio(self, audio_path: Path) -> torch.Tensor:
+        if not audio_path.is_file():
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+        waveform, sample_rate = torchaudio.load(str(audio_path))
+        if waveform.ndim != 2:
+            raise ValueError(f"Unsupported waveform shape: {tuple(waveform.shape)}")
+
+        waveform = waveform.mean(dim=0)
+        if sample_rate != self.sample_rate:
+            waveform = torchaudio.functional.resample(
+                waveform,
+                sample_rate,
+                self.sample_rate,
+            )
+        return waveform.unsqueeze(0).to(self.device)
+
+    @torch.inference_mode()
+    def transcribe(self, audio_path: Path) -> str:
+        wavs = self._load_audio(audio_path)
+        wav_lens = torch.ones(1, dtype=torch.float32, device=self.device)
+
+        features = self.compute_features(wavs)
+        features = self.normalize(features, wav_lens, epoch=0)
+        source = self.cnn(features)
+
+        dummy_bos = torch.full(
+            (wavs.shape[0], 1),
+            self.bos_index,
+            dtype=torch.long,
+            device=self.device,
+        )
+        encoder_output, _ = self.transformer(
+            source,
+            dummy_bos,
+            wav_lens,
+            pad_idx=self.pad_index,
+        )
+
+        search_output = self.search(encoder_output, wav_lens)
+        hypotheses = (
+            search_output[0]
+            if isinstance(search_output, (tuple, list))
+            else search_output
+        )
+        token_ids = hypotheses[0]
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+
+        return self.tokenizer.decode_ids(
+            [int(token) for token in token_ids]
+        ).strip()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Transcribe one audio file with ConMamba-small-ca."
+    )
+    parser.add_argument("audio", type=Path, help="Path to the input audio file")
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        help="PyTorch device used for inference (default: cuda)",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    recognizer = ConMambaASR(device=args.device)
+    print(recognizer.transcribe(args.audio.expanduser().resolve()))
+
+
+if __name__ == "__main__":
+    main()

tokenizer.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69808ba4897ce1a6ce492f0759a8326715106b60ca148842c362b9dc09463942
+size 253025