feat: add torchembed optional backend for categorical embeddings

pyproject.toml

@@ -75,6 +75,10 @@ extra = [
     "pytorch-tabnet<4.2",
 ]
 
+torchembed = [
+    "torchembed>=0.3.1",
+]
+
 notebooks = [
     "ipywidgets",
     "matplotlib>3.1",

src/pytorch_tabular/config/config.py

@@ -782,6 +782,12 @@ class ModelConfig:
 
         embedding_dropout (float): Dropout to be applied to the Categorical Embedding. Defaults to 0.0
 
+        embedding_backend (str): Backend to use for categorical embeddings. ``'native'`` (default) uses one
+                ``nn.Embedding`` per column. ``'torchembed'`` delegates to
+                ``torchembed.categorical.MultiCategoricalEmbedding``, which fuses all columns into a single
+                module with auto-sized dimensions. Requires ``pip install torchembed``.
+                Choices are: [``'native'``, ``'torchembed'``].
+
         batch_norm_continuous_input (bool): If True, we will normalize the continuous layer by passing it
                 through a BatchNorm layer.
 
@@ -851,6 +857,16 @@ class ModelConfig:
         default=0.0,
         metadata={"help": "Dropout to be applied to the Categorical Embedding. Defaults to 0.0"},
     )
+    embedding_backend: str = field(
+        default="native",
+        metadata={
+            "help": "Backend to use for categorical embeddings. 'native' (default) uses one nn.Embedding "
+            "per column. 'torchembed' delegates to torchembed.categorical.MultiCategoricalEmbedding, "
+            "which fuses all columns into a single module with auto-sized dimensions "
+            "(requires pip install torchembed).",
+            "choices": ["native", "torchembed"],
+        },
+    )
     batch_norm_continuous_input: bool = field(
         default=True,
         metadata={"help": "If True, we will normalize the continuous layer by passing it through a BatchNorm layer."},

src/pytorch_tabular/models/category_embedding/category_embedding_model.py

@@ -49,6 +49,7 @@ def _build_embedding_layer(self):
             embedding_dropout=self.hparams.embedding_dropout,
             batch_norm_continuous_input=self.hparams.batch_norm_continuous_input,
             virtual_batch_size=self.hparams.virtual_batch_size,
+            embedding_backend=getattr(self.hparams, "embedding_backend", "native"),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:

src/pytorch_tabular/models/common/layers/embeddings.py

@@ -109,7 +109,17 @@ def forward(self, x: Dict[str, Any]) -> torch.Tensor:
 
 
 class Embedding1dLayer(nn.Module):
-    """Enables different values in a categorical features to have different embeddings."""
+    """Enables different values in a categorical features to have different embeddings.
+
+    Supports two backends:
+
+    - ``"native"`` (default): uses one ``nn.Embedding`` per categorical column, exactly
+      as pytorch-tabular has always done.
+    - ``"torchembed"``: delegates to
+      ``torchembed.categorical.MultiCategoricalEmbedding``, which fuses all
+      per-column embeddings into a single module with auto-sized dimensions.
+      Requires the optional ``torchembed`` package (``pip install torchembed``).
+    """
 
     def __init__(
         self,
@@ -118,14 +128,35 @@ def __init__(
         embedding_dropout: float = 0.0,
         batch_norm_continuous_input: bool = False,
         virtual_batch_size: Optional[int] = None,
+        embedding_backend: str = "native",
     ):
         super().__init__()
         self.continuous_dim = continuous_dim
         self.categorical_embedding_dims = categorical_embedding_dims
         self.batch_norm_continuous_input = batch_norm_continuous_input
+        self.embedding_backend = embedding_backend
+
+        if embedding_backend == "torchembed":
+            try:
+                from torchembed.categorical import MultiCategoricalEmbedding
+            except ImportError as exc:
+                raise ImportError(
+                    "The 'torchembed' package is required when embedding_backend='torchembed'. "
+                    "Install it with: pip install torchembed"
+                ) from exc
+            cardinalities = [card for card, _ in categorical_embedding_dims]
+            self._torchembed_emb = MultiCategoricalEmbedding(cardinalities=cardinalities)
+            self._cat_output_dim: int = self._torchembed_emb.output_dim
+        elif embedding_backend == "native":
+            # Native per-column embedding layers (original behaviour)
+            self.cat_embedding_layers = nn.ModuleList([nn.Embedding(x, y) for x, y in categorical_embedding_dims])
+            self._cat_output_dim = sum(dim for _, dim in categorical_embedding_dims)
+        else:
+            raise ValueError(
+                f"Unknown embedding_backend '{embedding_backend}'. "
+                "Supported values are: 'native', 'torchembed'."
+            )
 
-        # Embedding layers
-        self.cat_embedding_layers = nn.ModuleList([nn.Embedding(x, y) for x, y in categorical_embedding_dims])
         if embedding_dropout > 0:
             self.embd_dropout = nn.Dropout(embedding_dropout)
         else:
@@ -134,16 +165,26 @@ def __init__(
         if batch_norm_continuous_input:
             self.normalizing_batch_norm = BatchNorm1d(continuous_dim, virtual_batch_size)
 
+    @property
+    def output_dim(self) -> int:
+        """Total output dimension of the categorical embeddings produced by this layer."""
+        return self._cat_output_dim
+
     def forward(self, x: Dict[str, Any]) -> torch.Tensor:
         assert "continuous" in x or "categorical" in x, "x must contain either continuous and categorical features"
         # (B, N)
         continuous_data, categorical_data = (
             x.get("continuous", torch.empty(0, 0)),
             x.get("categorical", torch.empty(0, 0)),
         )
-        assert categorical_data.shape[1] == len(
-            self.cat_embedding_layers
-        ), "categorical_data must have same number of columns as categorical embedding layers"
+        if self.embedding_backend == "torchembed":
+            assert categorical_data.shape[1] == len(
+                self.categorical_embedding_dims
+            ), "categorical_data must have same number of columns as categorical embedding dims"
+        else:
+            assert categorical_data.shape[1] == len(
+                self.cat_embedding_layers
+            ), "categorical_data must have same number of columns as categorical embedding layers"
         assert (
             continuous_data.shape[1] == self.continuous_dim
         ), "continuous_data must have same number of columns as continuous dim"
@@ -155,13 +196,17 @@ def forward(self, x: Dict[str, Any]) -> torch.Tensor:
                 embed = continuous_data
             # (B, N, C)
         if categorical_data.shape[1] > 0:
-            categorical_embed = torch.cat(
-                [
-                    embedding_layer(categorical_data[:, i])
-                    for i, embedding_layer in enumerate(self.cat_embedding_layers)
-                ],
-                dim=1,
-            )
+            if self.embedding_backend == "torchembed":
+                # torchembed returns (batch, output_dim) directly
+                categorical_embed = self._torchembed_emb(categorical_data)
+            else:
+                categorical_embed = torch.cat(
+                    [
+                        embedding_layer(categorical_data[:, i])
+                        for i, embedding_layer in enumerate(self.cat_embedding_layers)
+                    ],
+                    dim=1,
+                )
             # (B, N, C + C)
             if embed is None:
                 embed = categorical_embed

tests/test_torchembed_backend.py

@@ -0,0 +1,145 @@
+"""Tests for the torchembed embedding backend in Embedding1dLayer.
+
+These tests are deliberately self-contained: they exercise Embedding1dLayer
+directly without needing a full TabularModel training run, so they pass even
+in CI environments that have torchembed installed but no GPU.
+
+Run with:
+    pytest tests/test_torchembed_backend.py -v
+"""
+
+import pytest
+import torch
+
+from pytorch_tabular.models.common.layers.embeddings import Embedding1dLayer
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+CARDINALITIES = [50, 7, 120]  # vocab sizes for three categorical columns
+EMBEDDING_DIMS = [(card, min(50, (card + 1) // 2)) for card in CARDINALITIES]
+BATCH_SIZE = 16
+CONTINUOUS_DIM = 4
+
+
+def _make_batch(batch_size: int = BATCH_SIZE, n_cats: int = len(CARDINALITIES)):
+    """Return a minimal input dict compatible with Embedding1dLayer.forward."""
+    categorical = torch.stack(
+        [torch.randint(0, CARDINALITIES[i], (batch_size,)) for i in range(n_cats)],
+        dim=1,
+    )
+    continuous = torch.randn(batch_size, CONTINUOUS_DIM)
+    return {"categorical": categorical, "continuous": continuous}
+
+
+# ---------------------------------------------------------------------------
+# Native backend (always runs)
+# ---------------------------------------------------------------------------
+
+
+def test_native_backend_output_shape():
+    """Native backend concatenates per-column embeddings + continuous features."""
+    layer = Embedding1dLayer(
+        continuous_dim=CONTINUOUS_DIM,
+        categorical_embedding_dims=EMBEDDING_DIMS,
+        embedding_backend="native",
+    )
+    x = _make_batch()
+    out = layer(x)
+    expected_cat_dim = sum(dim for _, dim in EMBEDDING_DIMS)
+    assert out.shape == (BATCH_SIZE, CONTINUOUS_DIM + expected_cat_dim)
+
+
+def test_native_backend_output_dim_property():
+    layer = Embedding1dLayer(
+        continuous_dim=CONTINUOUS_DIM,
+        categorical_embedding_dims=EMBEDDING_DIMS,
+        embedding_backend="native",
+    )
+    expected = sum(dim for _, dim in EMBEDDING_DIMS)
+    assert layer.output_dim == expected
+
+
+def test_invalid_backend_raises():
+    with pytest.raises(ValueError, match="Unknown embedding_backend"):
+        Embedding1dLayer(
+            continuous_dim=CONTINUOUS_DIM,
+            categorical_embedding_dims=EMBEDDING_DIMS,
+            embedding_backend="nonexistent",
+        )
+
+
+# ---------------------------------------------------------------------------
+# torchembed backend (skipped when torchembed is not installed)
+# ---------------------------------------------------------------------------
+
+torchembed = pytest.importorskip(
+    "torchembed",
+    reason="torchembed not installed; skipping torchembed-backend tests. "
+    "Install with: pip install torchembed",
+)
+
+
+def test_torchembed_backend_output_shape():
+    """torchembed backend must produce the same 2-D output shape as the native backend."""
+    layer = Embedding1dLayer(
+        continuous_dim=CONTINUOUS_DIM,
+        categorical_embedding_dims=EMBEDDING_DIMS,
+        embedding_backend="torchembed",
+    )
+    x = _make_batch()
+    out = layer(x)
+
+    # output is 2-D: (batch, continuous_dim + cat_output_dim)
+    assert out.ndim == 2
+    assert out.shape[0] == BATCH_SIZE
+    assert out.shape[1] == CONTINUOUS_DIM + layer.output_dim
+
+
+def test_torchembed_backend_output_dim_property():
+    """output_dim must match MultiCategoricalEmbedding.output_dim."""
+    from torchembed.categorical import MultiCategoricalEmbedding
+
+    layer = Embedding1dLayer(
+        continuous_dim=CONTINUOUS_DIM,
+        categorical_embedding_dims=EMBEDDING_DIMS,
+        embedding_backend="torchembed",
+    )
+    ref = MultiCategoricalEmbedding(cardinalities=CARDINALITIES)
+    assert layer.output_dim == ref.output_dim
+
+
+def test_torchembed_backend_is_differentiable():
+    """Gradients must flow through the torchembed embedding layer."""
+    layer = Embedding1dLayer(
+        continuous_dim=CONTINUOUS_DIM,
+        categorical_embedding_dims=EMBEDDING_DIMS,
+        embedding_backend="torchembed",
+    )
+    x = _make_batch()
+    out = layer(x)
+    loss = out.sum()
+    loss.backward()
+    # If we reach here without RuntimeError, gradients are flowing fine.
+
+
+def test_torchembed_missing_import(monkeypatch):
+    """A helpful ImportError is raised when torchembed is not importable."""
+    import builtins
+
+    real_import = builtins.__import__
+
+    def mock_import(name, *args, **kwargs):
+        if name == "torchembed.categorical":
+            raise ImportError("mocked absence")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", mock_import)
+    with pytest.raises(ImportError, match="pip install torchembed"):
+        Embedding1dLayer(
+            continuous_dim=CONTINUOUS_DIM,
+            categorical_embedding_dims=EMBEDDING_DIMS,
+            embedding_backend="torchembed",
+        )