[ChrisW09]

Fixes #483

Problem

EmbeddingLayer.seq_len counted only numerical and categorical features, and the pooling architectures sized learned_flatten from their numerical, categorical and external-embedding features. Neither count included the CLS token, and seq_len also missed the external-embedding tokens. Two bugs followed:

• shuffle_embeddings=True (Mambular, MambAttention): x[:, perm, :] used a permutation that was too short. It always dropped the external-embedding tokens, so the model ignored fit(embeddings=...). With use_cls=True it dropped the last feature token and moved the CLS token away from position 0 (cls_position=0), or dropped the CLS token itself (cls_position=1).
• pooling_method="learned_flatten" with use_cls=True (FTTransformer, Mambular, MambAttention; TabTransformer and TabulaRNN share the sizing code but have no use_cls config field): the pooling layer was sized without the CLS token, so the first forward pass crashed with a shape mismatch.

Fix

The token count now has one source of truth in EmbeddingLayer:

• seq_len counts external-embedding features as well, since there is one token per feature.
• The new n_tokens attribute equals seq_len plus one when a CLS token is used, which is the sequence length forward() returns. Both are documented in the EmbeddingLayer docstring, and the custom-model guide sizes its example head from n_tokens.
• The new feature_permutation() draws the shuffle permutation over all feature tokens and keeps the CLS token at its configured position, so "cls" pooling still reads it.

The architectures use them as follows:

• Mambular and MambAttention register embedding_layer.feature_permutation() as their perm buffer.
• FTTransformer, Mambular, MambAttention, TabTransformer and TabulaRNN size their pooling layers with embedding_layer.n_tokens. SAINT and AutoInt already counted the CLS token and are unchanged.

Compatibility

Models without a CLS token or external embeddings get exactly the same permutation size and pooling layers as before, so their saved artifacts still load. A shuffled Mambular/MambAttention model saved by ≤ 2.1.0 with use_cls=True or external embeddings had a shorter perm buffer and silently dropped a token. Such a model now fails strict loading and has to be retrained. The model-zoo pages say so.

Tests

In tests/architectures/test_architecture_contracts.py:

• test_shuffled_embeddings_keep_every_token (Mambular, MambAttention; CLS first, CLS last, external embeddings): perm covers every token returned by the embedding layer and leaves the CLS token in place.
• test_shuffled_embeddings_use_external_embeddings: the output now changes when the external embeddings change.
• test_learned_flatten_pooling_counts_cls_token: the five architectures above run a forward pass with use_cls=True and learned_flatten pooling, with the CLS token in either position. TabTransformer and TabulaRNN get use_cls set on the config object, since their config classes have no such field.

18 of the 20 selected tests fail on main; the 2 that pass are the existing checkpoint tests. All 20 pass with this change.

Docs

• docs/model_zoo/stable/mambular.md and mambattention.md describe what the shuffle permutes and add the compatibility note above.
• The shuffle_embeddings docstrings in MambularConfig and MambAttentionConfig are updated.
• feature_permutation() has a NumPy-style docstring, and the EmbeddingLayer docstring documents emb_feature_info, seq_len and n_tokens.
• docs/core_concepts/custom_models.md sizes the example head from embedding_layer.n_tokens, so the example keeps working with a CLS token.

Checks

• just test: 3254 passed, 46 skipped, 28 xfailed; coverage 91.29%
• just docs: build succeeded (warnings treated as errors)
• just check: all hooks passed (ruff, prettier, pyright)