fix(nn): count CLS and embedding tokens in sequence lengths

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):

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):

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):

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):

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):

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

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

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__()
 
@@ -659,8 +670,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 +864,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"]

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)
 ```
 

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 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.
 
 ## Practical Config

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 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.
 
 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.

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