From 83508eeefaa7d605fc0edf8e351a0cada855f171 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 21:30:33 +0200
Subject: [PATCH 1/2] fix(architectures): apply the Mamba block in
 MambaTab.forward

self.mamba was constructed in __init__ but forward() went
initial_layer -> unsqueeze -> norm_f -> activation -> squeeze -> head and
never invoked it. Every MambaTab classifier/regressor/LSS silently
trained as a single linear layer with an MLP head, while all Mamba
parameters sat in the optimizer receiving no gradients.

Fixes #412

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 deeptab/architectures/mambatab.py |  1 +
 tests/test_mambatab_backbone.py   | 34 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 tests/test_mambatab_backbone.py

diff --git a/deeptab/architectures/mambatab.py b/deeptab/architectures/mambatab.py
index a70705ff..5aa75df4 100644
--- a/deeptab/architectures/mambatab.py
+++ b/deeptab/architectures/mambatab.py
@@ -110,6 +110,7 @@ def forward(self, *data):
 
         x = self.norm_f(x)
         x = self.embedding_activation(x)
+        x = self.mamba(x)
         if self.axis == 1:
             x = x.squeeze(1)
         else:
diff --git a/tests/test_mambatab_backbone.py b/tests/test_mambatab_backbone.py
new file mode 100644
index 00000000..2ff85e79
--- /dev/null
+++ b/tests/test_mambatab_backbone.py
@@ -0,0 +1,34 @@
+"""Regression test: MambaTab must actually run its Mamba block.
+
+forward() went initial_layer -> norm -> activation -> head and never called
+self.mamba, so the model trained as a linear layer with an MLP head while all
+Mamba parameters sat in the optimizer with no gradients.
+"""
+
+import numpy as np
+import pandas as pd
+
+from deeptab.models import MambaTabRegressor
+
+
+def _backward_one_batch(model):
+    task = model._task_model
+    task.zero_grad()
+    (num, cat, emb), labels = next(iter(model._data_module.train_dataloader()))
+    preds = task(num, cat, emb)
+    task.compute_loss(preds, labels).backward()
+    return task
+
+
+def test_mambatab_trains_its_mamba_block():
+    rng = np.random.RandomState(0)
+    X = pd.DataFrame({"num1": rng.randn(60), "num2": rng.rand(60)})
+    y = rng.randn(60)
+
+    model = MambaTabRegressor()
+    model.fit(X, y, max_epochs=1, batch_size=16, accelerator="cpu")
+
+    task = _backward_one_batch(model)
+    mamba_params = [(name, p) for name, p in task.named_parameters() if ".mamba." in name]
+    assert mamba_params, "expected mamba parameters on the estimator"
+    assert all(p.grad is not None for _, p in mamba_params)

From 82a3d574418101a8bce6d2ab8b0ae4f595e75dd2 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 22:00:23 +0200
Subject: [PATCH 2/2] test: satisfy pyright on the MambaTab backbone-gradient
 test

model._task_model and model._data_module are Optional-typed internals.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 tests/test_mambatab_backbone.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/test_mambatab_backbone.py b/tests/test_mambatab_backbone.py
index 2ff85e79..2eeeabb7 100644
--- a/tests/test_mambatab_backbone.py
+++ b/tests/test_mambatab_backbone.py
@@ -5,16 +5,18 @@
 Mamba parameters sat in the optimizer with no gradients.
 """
 
+from typing import Any, cast
+
 import numpy as np
 import pandas as pd
 
 from deeptab.models import MambaTabRegressor
 
 
-def _backward_one_batch(model):
-    task = model._task_model
+def _backward_one_batch(model) -> Any:
+    task = cast(Any, model._task_model)
     task.zero_grad()
-    (num, cat, emb), labels = next(iter(model._data_module.train_dataloader()))
+    (num, cat, emb), labels = next(iter(cast(Any, model._data_module).train_dataloader()))
     preds = task(num, cat, emb)
     task.compute_loss(preds, labels).backward()
     return task