From a724bdfd8a962405f1df62075343152e78885469 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Wed, 21 Oct 2026 15:02:53 +0200
Subject: [PATCH 1/2] fix(nn): count CLS and embedding tokens in sequence
 lengths

EmbeddingLayer.seq_len only counted numerical and categorical features,
so the shuffle_embeddings permutation of Mambular and MambAttention
dropped external-embedding tokens and a feature or the CLS token, and
learned_flatten pooling was sized without the CLS token. seq_len now
counts external-embedding features too, the new n_tokens attribute
matches the length returned by forward(), feature_permutation() shuffles
only feature tokens and keeps the CLS token in place, and the pooling
architectures size their pooling layers from n_tokens.

Fixes #483
---
 deeptab/architectures/ft_transformer.py       |  3 +-
 deeptab/architectures/mambattention.py        |  5 +-
 deeptab/architectures/mambular.py             |  5 +-
 deeptab/architectures/tabtransformer.py       |  9 ++-
 deeptab/architectures/tabularnn.py            |  3 +-
 .../configs/models/mambattention_config.py    |  3 +-
 deeptab/configs/models/mambular_config.py     |  3 +-
 deeptab/nn/blocks/common.py                   | 26 +++++++-
 docs/model_zoo/stable/mambattention.md        |  2 +
 docs/model_zoo/stable/mambular.md             |  2 +
 .../test_architecture_contracts.py            | 65 +++++++++++++++++++
 11 files changed, 107 insertions(+), 19 deletions(-)

diff --git a/deeptab/architectures/ft_transformer.py b/deeptab/architectures/ft_transformer.py
index 6fa3a0da..49c10975 100644
--- a/deeptab/architectures/ft_transformer.py
+++ b/deeptab/architectures/ft_transformer.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch.nn as nn
 
 from deeptab.core import BaseModel
@@ -85,7 +84,7 @@ def __init__(
         )
 
         # pooling
-        n_inputs = np.sum([len(info) for info in feature_information])
+        n_inputs = self.embedding_layer.n_tokens
         self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model)
 
     def forward(self, *data):
diff --git a/deeptab/architectures/mambattention.py b/deeptab/architectures/mambattention.py
index 709b01e6..2790d864 100644
--- a/deeptab/architectures/mambattention.py
+++ b/deeptab/architectures/mambattention.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 
 from deeptab.core import BaseModel
@@ -100,10 +99,10 @@ def __init__(
         )
 
         if self.shuffle_embeddings:
-            self.register_buffer("perm", torch.randperm(self.embedding_layer.seq_len))
+            self.register_buffer("perm", self.embedding_layer.feature_permutation())
 
         # pooling
-        n_inputs = np.sum([len(info) for info in feature_information])
+        n_inputs = self.embedding_layer.n_tokens
         self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model)
 
     def forward(self, *data):
diff --git a/deeptab/architectures/mambular.py b/deeptab/architectures/mambular.py
index be1bf593..c1cca7af 100644
--- a/deeptab/architectures/mambular.py
+++ b/deeptab/architectures/mambular.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 
 from deeptab.core import BaseModel
@@ -81,10 +80,10 @@ def __init__(
         )
 
         if self.hparams.shuffle_embeddings:
-            self.register_buffer("perm", torch.randperm(self.embedding_layer.seq_len))
+            self.register_buffer("perm", self.embedding_layer.feature_permutation())
 
         # pooling
-        n_inputs = np.sum([len(info) for info in feature_information])
+        n_inputs = self.embedding_layer.n_tokens
         self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model)
 
     def forward(self, *data):
diff --git a/deeptab/architectures/tabtransformer.py b/deeptab/architectures/tabtransformer.py
index fff4b299..0a7f3258 100644
--- a/deeptab/architectures/tabtransformer.py
+++ b/deeptab/architectures/tabtransformer.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 import torch.nn as nn
 
@@ -113,10 +112,10 @@ def __init__(
         )
 
         # pooling
-        # Only cat/embedding features pass through the encoder whose output
-        # gets pooled (numerical features are concatenated in separately),
-        # so n_inputs excludes num_feature_info.
-        n_inputs = int(np.sum([len(info) for info in feature_information[1:]]))
+        # Only cat/embedding features (plus the CLS token, if enabled) pass
+        # through the encoder whose output gets pooled (numerical features are
+        # concatenated in separately), so n_inputs excludes num_feature_info.
+        n_inputs = self.embedding_layer.n_tokens
         self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model)
 
     def forward(self, *data):
diff --git a/deeptab/architectures/tabularnn.py b/deeptab/architectures/tabularnn.py
index da8dea7e..0304628f 100644
--- a/deeptab/architectures/tabularnn.py
+++ b/deeptab/architectures/tabularnn.py
@@ -1,6 +1,5 @@
 from dataclasses import replace
 
-import numpy as np
 import torch
 import torch.nn as nn
 
@@ -84,7 +83,7 @@ def __init__(
         self.norm_f = get_normalization_layer(temp_config)
 
         # pooling
-        n_inputs = int(np.sum([len(info) for info in feature_information]))
+        n_inputs = self.embedding_layer.n_tokens
         self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.dim_feedforward)
 
     def forward(self, *data):
diff --git a/deeptab/configs/models/mambattention_config.py b/deeptab/configs/models/mambattention_config.py
index c816b44e..f3524c9c 100644
--- a/deeptab/configs/models/mambattention_config.py
+++ b/deeptab/configs/models/mambattention_config.py
@@ -57,7 +57,8 @@ class MambAttentionConfig(BaseModelConfig):
     BC_layer_norm : bool, default=False
         Whether to apply layer normalization to B-C matrices.
     shuffle_embeddings : bool, default=False
-        Whether to shuffle embeddings before passing to Mamba layers.
+        Whether to shuffle the feature tokens before passing them to the Mamba layers. The CLS token,
+        if used, keeps its position.
     head_layer_sizes : list, default=[]
         Sizes of the fully connected layers in the model's head.
     head_dropout : float, default=0.5
diff --git a/deeptab/configs/models/mambular_config.py b/deeptab/configs/models/mambular_config.py
index 1a978c9f..2ddf9d3e 100644
--- a/deeptab/configs/models/mambular_config.py
+++ b/deeptab/configs/models/mambular_config.py
@@ -51,7 +51,8 @@ class MambularConfig(BaseModelConfig):
     BC_layer_norm : bool, default=False
         Whether to use layer norm on the B and C matrices
     shuffle_embeddings : bool, default=False
-        Whether to shuffle embeddings before being passed to Mamba layers.
+        Whether to shuffle the feature tokens before passing them to the Mamba layers. The CLS token,
+        if used, keeps its position.
     head_layer_sizes : list, default=[]
         Sizes of the layers in the model's head.
     head_dropout : float, default=0.5
diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py
index 48fcb058..4af005a8 100644
--- a/deeptab/nn/blocks/common.py
+++ b/deeptab/nn/blocks/common.py
@@ -659,8 +659,11 @@ def __init__(self, num_feature_info, cat_feature_info, emb_feature_info, config)
         self.embedding_bias = getattr(config, "embedding_bias", False)
         self.cat_encoding = getattr(config, "cat_encoding", "int")
 
-        # Sequence length
-        self.seq_len = len(num_feature_info) + len(cat_feature_info)
+        # Sequence length: one token per numerical, categorical, and external
+        # embedding feature. n_tokens additionally counts the CLS token, so it
+        # matches the sequence length returned by forward().
+        self.seq_len = len(num_feature_info) + len(cat_feature_info) + len(emb_feature_info)
+        self.n_tokens = self.seq_len + int(self.use_cls)
 
         # Initialize numerical embeddings based on embedding_type
         if self.embedding_type == "ndt":
@@ -850,6 +853,25 @@ def forward(self, num_features, cat_features, emb_features):
 
         return x
 
+    def feature_permutation(self):
+        """Draw a random permutation of the feature tokens returned by :meth:`forward`.
+
+        Every feature token (numerical, categorical, and external embedding
+        features) is shuffled, while the CLS token, when enabled, keeps its
+        configured position so that ``"cls"`` pooling still finds it.
+
+        Returns
+        -------
+        Tensor
+            Index tensor of length ``n_tokens`` for the sequence dimension.
+        """
+        perm = torch.randperm(self.seq_len)
+        if not self.use_cls:
+            return perm
+        if self.cls_position == 0:
+            return torch.cat([torch.zeros(1, dtype=perm.dtype), perm + 1])
+        return torch.cat([perm, torch.full((1,), self.seq_len, dtype=perm.dtype)])
+
     def check_plr_embedding_compatibility(self, feature_info: tuple):
         # List of incompatible preprocessing terms for PLR embedding
         incompatible_terms = ["ple", "one-hot", "polynomial", "splines", "sigmoid", "rbf"]
diff --git a/docs/model_zoo/stable/mambattention.md b/docs/model_zoo/stable/mambattention.md
index a5611e3f..3d0be31a 100644
--- a/docs/model_zoo/stable/mambattention.md
+++ b/docs/model_zoo/stable/mambattention.md
@@ -40,6 +40,8 @@ The default config uses `d_model=64`, `n_layers=4`, `n_heads=8`, `n_attention_la
 
 With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them.
 
+The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 permuted only the numerical and categorical tokens, which silently dropped the CLS token or the external-embedding tokens. Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
+
 > **Configuration note:** Pooling operates on width-`d_model` tokens, not TabulaRNN's `dim_feedforward` width. The local Mamba residual blocks also use the shared fixes for dilated convolution padding, missing-`mambapy` sequential fallback, and `AD_weight_decay=False` parameter markers. Final normalization after pooling is selected by `norm` and is separate from normalization inside the residual blocks.
 
 ## Practical Config
diff --git a/docs/model_zoo/stable/mambular.md b/docs/model_zoo/stable/mambular.md
index fdb002f9..c14f12cb 100644
--- a/docs/model_zoo/stable/mambular.md
+++ b/docs/model_zoo/stable/mambular.md
@@ -39,6 +39,8 @@ The default config uses `d_model=64`, `n_layers=4`, `d_state=128`, `d_conv=4`, `
 
 With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them.
 
+The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 permuted only the numerical and categorical tokens, which silently dropped the CLS token or the external-embedding tokens. Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
+
 > **Configuration note:** In the local Mamba block, `dilation=3` with `d_conv=3` gives padding 6 in each directional convolution. Padding and the existing crop preserve sequence length; previously, the forward convolution ignored dilation. If `use_pscan=True` but `mambapy` cannot be imported, the block switches to its sequential scan rather than calling a missing function. `AD_weight_decay=False` marks `A_log` and `D` for zero weight decay; module-based optimizer construction respects those markers independently of the bias/norm exemption flag.
 
 Learnable pooling uses the sequence width `d_model`, not a `dim_feedforward` field. For `d_model=16`, attention pooling learns a width-16 weight vector. This keeps the shared pooling helper independent of TabulaRNN's recurrent width.
diff --git a/tests/architectures/test_architecture_contracts.py b/tests/architectures/test_architecture_contracts.py
index ec9090b5..4020cce6 100644
--- a/tests/architectures/test_architecture_contracts.py
+++ b/tests/architectures/test_architecture_contracts.py
@@ -199,6 +199,71 @@ def test_shuffled_embeddings_permutation_is_checkpointed(architecture_name):
     torch.testing.assert_close(restored(features, [], []), model(features, [], []), rtol=0, atol=0)
 
 
+def _small_sequence_config(architecture_name, **options):
+    model_class = next(model_class for model_class in model_classes if model_class.__name__ == architecture_name)
+    config = get_model_config(model_class)
+    for name, value in {"d_model": 8, "n_layers": 1, "d_state": 4, "n_heads": 2, "dim_feedforward": 8}.items():
+        if hasattr(config, name):
+            setattr(config, name, value)
+    for name, value in options.items():
+        setattr(config, name, value)
+    return model_class, config
+
+
+_NUM_AND_EMB_INFO = (
+    {f"value_{index}": {"dimension": 1} for index in range(3)},
+    {},
+    {"text": {"dimension": 5}},
+)
+
+
+@pytest.mark.parametrize("architecture_name", ["Mambular", "MambAttention"])
+@pytest.mark.parametrize(
+    "options, feature_info",
+    [
+        pytest.param({"use_cls": True, "cls_position": 0}, (*_NUM_AND_EMB_INFO[:2], {}), id="cls-first"),
+        pytest.param({"use_cls": True, "cls_position": 1}, (*_NUM_AND_EMB_INFO[:2], {}), id="cls-last"),
+        pytest.param({}, _NUM_AND_EMB_INFO, id="external-embeddings"),
+    ],
+)
+def test_shuffled_embeddings_keep_every_token(architecture_name, options, feature_info):
+    model_class, config = _small_sequence_config(architecture_name, shuffle_embeddings=True, **options)
+    model = model_class(feature_info, config=config).eval()
+    embeddings = [torch.randn(4, 5)] if feature_info[2] else []
+    tokens = model.embedding_layer([torch.randn(4, 1) for _ in range(3)], [], embeddings)
+
+    assert model.embedding_layer.n_tokens == tokens.shape[1]
+    assert sorted(model.perm.tolist()) == list(range(tokens.shape[1]))
+    if options.get("use_cls"):
+        cls_index = 0 if options["cls_position"] == 0 else tokens.shape[1] - 1
+        assert model.perm[cls_index].item() == cls_index
+
+
+@pytest.mark.parametrize("architecture_name", ["Mambular", "MambAttention"])
+def test_shuffled_embeddings_use_external_embeddings(architecture_name):
+    model_class, config = _small_sequence_config(architecture_name, shuffle_embeddings=True)
+    model = model_class(_NUM_AND_EMB_INFO, config=config).eval()
+    features = [torch.randn(4, 1) for _ in range(3)]
+    with torch.no_grad():
+        first = model(features, [], [torch.randn(4, 5)])
+        second = model(features, [], [torch.randn(4, 5)])
+    assert not torch.allclose(first, second)
+
+
+@pytest.mark.parametrize(
+    "architecture_name", ["FTTransformer", "Mambular", "MambAttention", "TabTransformer", "TabulaRNN"]
+)
+@pytest.mark.parametrize("cls_position", [0, 1])
+def test_learned_flatten_pooling_counts_cls_token(architecture_name, cls_position):
+    model_class, config = _small_sequence_config(
+        architecture_name, use_cls=True, cls_position=cls_position, pooling_method="learned_flatten"
+    )
+    feature_info = ({"value": {"dimension": 1}}, {"color": {"dimension": 1, "categories": 3}}, {})
+    model = model_class(feature_info, num_classes=2, config=config)
+    output = model([torch.randn(4, 1)], [torch.randint(0, 3, (4, 1))], [])
+    assert output.shape == (4, 2)
+
+
 @pytest.mark.parametrize("model_type", ["mLSTM", "sLSTM"])
 def test_tabularnn_recurrent_options_preserve_rows(model_type):
     from deeptab.architectures.tabularnn import TabulaRNN

From 4f5a62df6d65edfcf476f2bdf1b06a716d45cf4d Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Wed, 21 Oct 2026 23:51:32 +0200
Subject: [PATCH 2/2] docs: describe the dropped tokens precisely and document
 n_tokens

With cls_position=0 the old permutation dropped the last feature token
and moved the CLS token, not the CLS token itself. EmbeddingLayer now
documents seq_len and n_tokens, and the custom-model guide sizes its
head from n_tokens.

Refs #483
---
 deeptab/nn/blocks/common.py            | 11 +++++++++++
 docs/core_concepts/custom_models.md    |  8 ++++----
 docs/model_zoo/stable/mambattention.md |  2 +-
 docs/model_zoo/stable/mambular.md      |  2 +-
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py
index 4af005a8..21595d8a 100644
--- a/deeptab/nn/blocks/common.py
+++ b/deeptab/nn/blocks/common.py
@@ -639,8 +639,19 @@ def __init__(self, num_feature_info, cat_feature_info, emb_feature_info, config)
         cat_feature_info : dict
             Dictionary where keys are categorical feature names and values are the number of categories
             for each feature.
+        emb_feature_info : dict
+            Dictionary where keys are external embedding feature names and values describe their
+            dimensions.
         config : Config
             Configuration object containing all required settings.
+
+        Attributes
+        ----------
+        seq_len : int
+            Number of feature tokens: one per numerical, categorical and external embedding feature.
+        n_tokens : int
+            Length of the token sequence returned by :meth:`forward`: ``seq_len``, plus one when
+            ``use_cls`` adds a CLS token.
         """
         super().__init__()
 
diff --git a/docs/core_concepts/custom_models.md b/docs/core_concepts/custom_models.md
index de78adcd..c1ac7d05 100644
--- a/docs/core_concepts/custom_models.md
+++ b/docs/core_concepts/custom_models.md
@@ -200,14 +200,14 @@ class MyEmbeddedModel(BaseModel):
         self.save_hyperparameters(ignore=["feature_information"])
 
         self.embedding_layer = EmbeddingLayer(*feature_information, config=config)
-        n_features = sum(len(info) for info in feature_information)
-        input_dim = n_features * self.hparams.d_model
+        # n_tokens counts every token the layer returns, including a CLS token if enabled.
+        input_dim = self.embedding_layer.n_tokens * self.hparams.d_model
 
         self.head = nn.Linear(input_dim, num_classes)
 
     def forward(self, *data):
-        x = self.embedding_layer(*data)   # (batch, n_features, d_model)
-        x = x.reshape(x.shape[0], -1)     # flatten to (batch, n_features * d_model)
+        x = self.embedding_layer(*data)   # (batch, n_tokens, d_model)
+        x = x.reshape(x.shape[0], -1)     # flatten to (batch, n_tokens * d_model)
         return self.head(x)
 ```
 
diff --git a/docs/model_zoo/stable/mambattention.md b/docs/model_zoo/stable/mambattention.md
index 3d0be31a..c2986ded 100644
--- a/docs/model_zoo/stable/mambattention.md
+++ b/docs/model_zoo/stable/mambattention.md
@@ -40,7 +40,7 @@ The default config uses `d_model=64`, `n_layers=4`, `n_heads=8`, `n_attention_la
 
 With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them.
 
-The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 permuted only the numerical and categorical tokens, which silently dropped the CLS token or the external-embedding tokens. Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
+The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 drew the permutation over the numerical and categorical tokens only, so the remaining positions were dropped: every external-embedding token and, with `use_cls=True`, either the last feature token (`cls_position=0`, which also moved the CLS token away from position 0) or the CLS token itself (`cls_position=1`). Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
 
 > **Configuration note:** Pooling operates on width-`d_model` tokens, not TabulaRNN's `dim_feedforward` width. The local Mamba residual blocks also use the shared fixes for dilated convolution padding, missing-`mambapy` sequential fallback, and `AD_weight_decay=False` parameter markers. Final normalization after pooling is selected by `norm` and is separate from normalization inside the residual blocks.
 
diff --git a/docs/model_zoo/stable/mambular.md b/docs/model_zoo/stable/mambular.md
index c14f12cb..024ebe56 100644
--- a/docs/model_zoo/stable/mambular.md
+++ b/docs/model_zoo/stable/mambular.md
@@ -39,7 +39,7 @@ The default config uses `d_model=64`, `n_layers=4`, `d_state=128`, `d_conv=4`, `
 
 With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them.
 
-The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 permuted only the numerical and categorical tokens, which silently dropped the CLS token or the external-embedding tokens. Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
+The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 drew the permutation over the numerical and categorical tokens only, so the remaining positions were dropped: every external-embedding token and, with `use_cls=True`, either the last feature token (`cls_position=0`, which also moved the CLS token away from position 0) or the CLS token itself (`cls_position=1`). Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them.
 
 > **Configuration note:** In the local Mamba block, `dilation=3` with `d_conv=3` gives padding 6 in each directional convolution. Padding and the existing crop preserve sequence length; previously, the forward convolution ignored dilation. If `use_pscan=True` but `mambapy` cannot be imported, the block switches to its sequential scan rather than calling a missing function. `AD_weight_decay=False` marks `A_log` and `D` for zero weight decay; module-based optimizer construction respects those markers independently of the bias/norm exemption flag.