Add NVIDIA RE-USE (SEMamba) speech enhancement model

src/mobius/_registry.py

@@ -113,6 +113,8 @@
     Qwen35VL3ModelCausalLMModel,
     Qwen35VLTextModel,
     QwenCausalLMModel,
+    ReUseConfig,
+    SEMambaSpeechEnhancementModel,
     SmolLM3CausalLMModel,
     SortformerDiarizationModel,
     WhisperForConditionalGeneration,
@@ -857,6 +859,16 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     ),
     "fastconformer_rnnt": ModelRegistration(EncDecRNNTModel, task="fastconformer-rnnt"),
     "sortformer": ModelRegistration(SortformerDiarizationModel, task="diarization"),
+    "reuse": ModelRegistration(
+        SEMambaSpeechEnhancementModel,
+        task="speech-enhancement",
+        config_class=ReUseConfig,
+    ),
+    "semamba": ModelRegistration(
+        SEMambaSpeechEnhancementModel,
+        task="speech-enhancement",
+        config_class=ReUseConfig,
+    ),
 }
 
 

src/mobius/components/__init__.py

@@ -68,6 +68,8 @@
     "RmsNorm2d",
     "ScaleFreeRMSNorm",
     "SelectiveScan",
+    "SequenceMambaBlock",
+    "SequenceSelectiveScan",
     "SiLU",
     "Siglip2NaFlexVisionEmbeddings",
     "Siglip2NaFlexVisionModel",
@@ -184,6 +186,7 @@
 from mobius.components._lora import LoRALinear
 from mobius.components._mamba_block import Mamba2Block as Mamba2Block
 from mobius.components._mamba_block import MambaBlock as MambaBlock
+from mobius.components._mamba_block import SequenceMambaBlock
 from mobius.components._mlp import FCMLP, MLP, FusedGateUpMLP, GatedMLP
 from mobius.components._mobilenetv5 import MobileNetV5Encoder
 from mobius.components._moe import (
@@ -306,6 +309,7 @@
 )
 from mobius.components._ssm import (
     SelectiveScan,
+    SequenceSelectiveScan,
 )
 from mobius.components._vision import (
     PatchEmbedding,

src/mobius/components/_conv.py

@@ -77,6 +77,7 @@ def __init__(
         stride: int | tuple[int, int] = 1,
         padding: int | tuple[int, int, int, int] = 0,
         groups: int = 1,
+        dilation: int | tuple[int, int] = 1,
     ):
         super().__init__()
         kernel_h, kernel_w = _pair(kernel_size, "kernel_size")
@@ -87,6 +88,7 @@ def __init__(
         self._strides = _pair(stride, "stride")
         self._pads = _resolve_pads(padding)
         self._groups = groups
+        self._dilations = _pair(dilation, "dilation")
 
     def forward(self, op: OpBuilder, x: ir.Value):
         return op.Conv(
@@ -96,6 +98,7 @@ def forward(self, op: OpBuilder, x: ir.Value):
             kernel_shape=list(self._kernel_size),
             strides=list(self._strides),
             pads=self._pads,
+            dilations=list(self._dilations),
             group=self._groups,
         )
 

src/mobius/components/_conv_test.py

@@ -113,6 +113,33 @@ def test_asymmetric_kernel_stride_and_padding(self):
         assert conv._strides == (1, 4)
         assert conv._pads == [0, 2, 0, 2]
 
+    def test_dilation_defaults_to_one(self):
+        conv = Conv2d(3, 16, kernel_size=3, padding=1)
+        assert conv._dilations == (1, 1)
+
+    def test_dilation_accepts_int_and_pair(self):
+        assert Conv2d(3, 16, kernel_size=3, dilation=2)._dilations == (2, 2)
+        assert Conv2d(3, 16, kernel_size=3, dilation=(4, 1))._dilations == (4, 1)
+
+    def test_positional_call_is_unchanged_by_dilation(self):
+        """``dilation`` was appended, so existing positional calls still work."""
+        conv = Conv2d(3, 16, 3, 2, 1, 1)
+        assert list(conv.weight.shape) == [16, 3, 3, 3]
+        assert conv._strides == (2, 2)
+        assert conv._pads == [1, 1, 1, 1]
+        assert conv._groups == 1
+        assert conv._dilations == (1, 1)
+
+    def test_dilation_reaches_the_conv_node(self):
+        conv = Conv2d(4, 4, kernel_size=(3, 3), padding=(2, 1, 2, 1), dilation=(2, 1))
+        test_builder, op, graph = create_test_builder()
+        x = create_test_input(test_builder, "x", [1, 4, 8, 8])
+
+        conv(op, x)
+
+        node = next(n for n in graph if n.op_type == "Conv")
+        assert list(node.attributes["dilations"].value) == [2, 1]
+
 
 class TestConv2dNoBias:
     """Tests for 2D convolution without bias."""

src/mobius/components/_mamba_block.py

@@ -24,7 +24,7 @@
 
 from mobius.components._common import INT64_MAX, Linear
 from mobius.components._rms_norm import GatedRMSNorm
-from mobius.components._ssm import SelectiveScan
+from mobius.components._ssm import SelectiveScan, SequenceSelectiveScan
 
 
 class _DepthwiseConv1d(nn.Module):
@@ -171,6 +171,84 @@ def forward(
         return output, new_conv_state, new_ssm_state
 
 
+class SequenceMambaBlock(nn.Module):
+    """Mamba1 layer applied to a whole sequence, with no carried state.
+
+    Same parameters and math as :class:`MambaBlock` (input projection →
+    causal Conv1D → selective scan → SiLU gate → output projection), but the
+    entire ``(batch, seq_len, d_model)`` sequence is processed in one call
+    and both the conv and SSM states start at zero.  This mirrors the
+    reference ``mamba_ssm.Mamba`` module's offline forward pass and is what
+    non-autoregressive backbones (e.g. speech enhancement) need.
+
+    Args:
+        d_model: Model hidden dimension.
+        d_inner: Expanded inner dimension (typically ``expand * d_model``).
+        d_state: SSM state dimension (typically 16).
+        dt_rank: Rank of the SSM time-step projection.
+            Defaults to ``ceil(d_model / 16)`` (Mamba convention).
+        conv_kernel: Causal Conv1D kernel size (typically 4).
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_inner: int,
+        d_state: int = 16,
+        dt_rank: int | None = None,
+        conv_kernel: int = 4,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.d_inner = d_inner
+        self.d_state = d_state
+        self.conv_kernel = conv_kernel
+        self.dt_rank = dt_rank if dt_rank is not None else -(-d_model // 16)
+
+        self.in_proj = Linear(d_model, 2 * d_inner, bias=False)
+        self.conv1d = _DepthwiseConv1d(d_inner, conv_kernel, bias=True)
+        self.ssm = SequenceSelectiveScan(d_inner, d_state, self.dt_rank)
+        self.out_proj = Linear(d_inner, d_model, bias=False)
+
+    def forward(self, op: OpBuilder, hidden_states: ir.Value):
+        """Run the Mamba layer over a full sequence.
+
+        Args:
+            op: ONNX op builder.
+            hidden_states: (batch, seq_len, d_model).
+
+        Returns:
+            (batch, seq_len, d_model) output.
+        """
+        # --- Step 1: Input projection, split into SSM branch and gate ---
+        projected = self.in_proj(op, hidden_states)  # (batch, seq_len, 2*d_inner)
+        x_branch, z_gate = op.Split(
+            projected,
+            [self.d_inner, self.d_inner],
+            axis=-1,
+            _outputs=2,
+        )
+
+        # --- Step 2: Causal depthwise Conv1D over the sequence ---
+        # (batch, seq_len, d_inner) → (batch, d_inner, seq_len) for Conv.
+        x_t = op.Transpose(x_branch, perm=[0, 2, 1])
+        # Left-pad by (conv_kernel - 1) so output position t only sees inputs
+        # up to t; the conv itself uses pads=[0, 0] and keeps seq_len.
+        x_padded = op.Pad(
+            x_t,
+            op.Constant(value_ints=[0, 0, self.conv_kernel - 1, 0, 0, 0]),
+        )
+        conv_out = op.Swish(self.conv1d(op, x_padded))  # (batch, d_inner, seq_len)
+
+        # --- Step 3: Selective scan over the sequence ---
+        x_ssm = op.Transpose(conv_out, perm=[0, 2, 1])  # (batch, seq_len, d_inner)
+        y = self.ssm(op, x_ssm)  # (batch, seq_len, d_inner)
+
+        # --- Step 4: Gating and output projection ---
+        gated = op.Mul(y, op.Swish(z_gate))
+        return self.out_proj(op, gated)
+
+
 # =====================================================================
 # Mamba2 block using LinearAttention
 # =====================================================================

src/mobius/components/_mamba_block_test.py

@@ -5,12 +5,18 @@
 
 from __future__ import annotations
 
+import numpy as np
+import onnx_ir as ir
+import pytest
+from onnxscript import GraphBuilder
+
+from mobius._constants import OPSET_VERSION
 from mobius._testing import (
     count_op_type,
     create_test_builder,
     create_test_input,
 )
-from mobius.components._mamba_block import MambaBlock
+from mobius.components._mamba_block import MambaBlock, SequenceMambaBlock
 
 
 class TestMambaBlock:
@@ -131,3 +137,116 @@ def test_custom_conv_kernel(self):
         block = MambaBlock(d_model=32, d_inner=64, conv_kernel=8)
         assert list(block.conv1d.weight.shape) == [64, 1, 8]
         assert block.conv_kernel == 8
+
+
+class TestSequenceMambaBlock:
+    """Tests for the full-sequence (stateless) Mamba1 block."""
+
+    def test_parameters_match_decode_block(self):
+        """Parameter names and shapes are identical to MambaBlock's."""
+
+        def spec(module):
+            return {name: list(p.shape) for name, p in module.named_parameters()}
+
+        assert spec(SequenceMambaBlock(d_model=32, d_inner=64, d_state=8)) == spec(
+            MambaBlock(d_model=32, d_inner=64, d_state=8)
+        )
+
+    def test_forward_needs_no_state(self):
+        """The whole sequence is consumed in one call, with no carried state."""
+        block = SequenceMambaBlock(d_model=32, d_inner=64, d_state=8)
+        test_builder, op, graph = create_test_builder()
+        hidden = create_test_input(test_builder, "hidden_states", [2, 7, 32])
+
+        out = block(op, hidden)
+
+        assert out is not None
+        assert count_op_type(graph, "Scan") == 1
+        # in_proj, x_proj, dt_proj, out_proj
+        assert count_op_type(graph, "MatMul") >= 4
+
+    def test_conv_is_left_padded_for_causality(self):
+        """Conv1D is left-padded by kernel_size - 1 so it stays causal."""
+        block = SequenceMambaBlock(d_model=16, d_inner=32, d_state=4, conv_kernel=4)
+        test_builder, op, graph = create_test_builder()
+        hidden = create_test_input(test_builder, "hidden_states", [1, 5, 16])
+
+        block(op, hidden)
+
+        pad = next(node for node in graph if node.op_type == "Pad")
+        pads = pad.inputs[1].const_value.numpy().tolist()
+        # (batch, d_inner, seq): 3 begin values then 3 end values.
+        assert pads == [0, 0, 3, 0, 0, 0]
+
+    def test_matches_reference_recurrence(self):
+        """ONNX output matches an explicit PyTorch selective-scan reference."""
+        pytest.importorskip("onnxruntime")
+        import onnxruntime as ort
+        import torch
+        import torch.nn.functional as torch_f
+
+        d_model, d_state, conv_kernel, expand = 8, 4, 4, 2
+        d_inner = d_model * expand
+        dt_rank = -(-d_model // 16)
+        batch, seq_len = 2, 6
+
+        torch.manual_seed(0)
+        weights = {
+            "in_proj.weight": torch.randn(2 * d_inner, d_model) * 0.2,
+            "conv1d.weight": torch.randn(d_inner, 1, conv_kernel) * 0.2,
+            "conv1d.bias": torch.randn(d_inner) * 0.2,
+            "ssm.x_proj.weight": torch.randn(dt_rank + 2 * d_state, d_inner) * 0.2,
+            "ssm.dt_proj.weight": torch.randn(d_inner, dt_rank) * 0.2,
+            "ssm.dt_proj.bias": torch.randn(d_inner) * 0.2,
+            "ssm.A_log": torch.randn(d_inner, d_state) * 0.2,
+            "ssm.D": torch.randn(d_inner) * 0.2,
+            "out_proj.weight": torch.randn(d_model, d_inner) * 0.2,
+        }
+
+        def reference(u):
+            """``mamba_ssm.Mamba`` forward, written out with plain torch ops."""
+            x, z = (u @ weights["in_proj.weight"].t()).chunk(2, dim=-1)
+            x = torch_f.conv1d(
+                torch_f.pad(x.transpose(1, 2), (conv_kernel - 1, 0)),
+                weights["conv1d.weight"],
+                weights["conv1d.bias"],
+                groups=d_inner,
+            )
+            x = torch_f.silu(x).transpose(1, 2)
+
+            x_dbl = x @ weights["ssm.x_proj.weight"].t()
+            dt_raw, b_mat, c_mat = torch.split(x_dbl, [dt_rank, d_state, d_state], dim=-1)
+            dt = torch_f.softplus(
+                dt_raw @ weights["ssm.dt_proj.weight"].t() + weights["ssm.dt_proj.bias"]
+            )
+
+            a_neg = -torch.exp(weights["ssm.A_log"])
+            state = torch.zeros(u.shape[0], d_inner, d_state)
+            outputs = []
+            for t in range(u.shape[1]):
+                dt_col = dt[:, t].unsqueeze(-1)
+                decay = torch.exp(dt_col * a_neg.unsqueeze(0))
+                update = dt_col * x[:, t].unsqueeze(-1) * b_mat[:, t].unsqueeze(1)
+                state = decay * state + update
+                outputs.append((state * c_mat[:, t].unsqueeze(1)).sum(-1))
+            y = torch.stack(outputs, dim=1) + weights["ssm.D"] * x
+            return (y * torch_f.silu(z)) @ weights["out_proj.weight"].t()
+
+        graph = ir.Graph([], [], nodes=[], name="g", opset_imports={"": OPSET_VERSION})
+        builder = GraphBuilder(graph)
+        hidden = builder.input(
+            "hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "seq", d_model]
+        )
+        block = SequenceMambaBlock(d_model, d_inner, d_state, dt_rank, conv_kernel)
+        for name, param in block.named_parameters():
+            param.const_value = ir.tensor(weights[name].numpy())
+        builder.add_output(block(builder.op, hidden), "y")
+
+        session = ort.InferenceSession(
+            ir.to_proto(ir.Model(graph, ir_version=11)).SerializeToString(),
+            providers=["CPUExecutionProvider"],
+        )
+        hidden_states = torch.randn(batch, seq_len, d_model)
+        (got,) = session.run(None, {"hidden_states": hidden_states.numpy()})
+
+        np.testing.assert_allclose(got, reference(hidden_states).detach().numpy(), atol=1e-5)