feat: add torchembed optional backend for categorical embeddings - #663
Open
py-ai-dev wants to merge 1 commit into
Open
feat: add torchembed optional backend for categorical embeddings#663py-ai-dev wants to merge 1 commit into
py-ai-dev wants to merge 1 commit into
Conversation
Adds an `embedding_backend` config option to `ModelConfig` (default `"native"`) that accepts `"torchembed"` as a value. When selected, `Embedding1dLayer` delegates to `torchembed.categorical.MultiCategoricalEmbedding` instead of building one `nn.Embedding` per column. Benefits over the native path: - Auto-sized embedding dimensions (standard `min(50,(n+1)//2)` rule) computed and managed inside a single fused module - Cleaner API: one module owns all categorical columns - Drop-in: output shape is identical so no downstream changes are needed Changes: - `Embedding1dLayer` gains `embedding_backend` param and an `output_dim` property (useful for both backends) - `ModelConfig` gains the `embedding_backend` field with docstring - `CategoryEmbeddingModel._build_embedding_layer` forwards the option - `pyproject.toml` gains a `[torchembed]` optional-dependency group - `tests/test_torchembed_backend.py` covers native path, torchembed output shapes, differentiability, and helpful ImportError messaging Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds an optional categorical-embedding backend (torchembed) to pytorch-tabular so users can switch embedding implementations via config while keeping the same downstream tensor shapes.
Changes:
- Introduces
embedding_backend(default"native") onModelConfig, and wires it throughCategoryEmbeddingModel. - Extends
Embedding1dLayerto support atorchembed-powered path and adds anoutput_dimproperty. - Adds an optional dependency group for
torchembedand a new unit test module covering both backends.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_torchembed_backend.py | Adds backend-focused unit tests (but currently skips the whole module when torchembed is absent). |
| src/pytorch_tabular/models/common/layers/embeddings.py | Implements embedding_backend switch and output_dim; needs a guard for cat_embedding_layers access under torchembed. |
| src/pytorch_tabular/models/category_embedding/category_embedding_model.py | Forwards embedding_backend from hparams into Embedding1dLayer. |
| src/pytorch_tabular/config/config.py | Adds embedding_backend field and help text to ModelConfig. |
| pyproject.toml | Adds torchembed optional-dependency group. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+78
to
+82
| torchembed = pytest.importorskip( | ||
| "torchembed", | ||
| reason="torchembed not installed; skipping torchembed-backend tests. " | ||
| "Install with: pip install torchembed", | ||
| ) |
Comment on lines
+102
to
+104
| """output_dim must match MultiCategoricalEmbedding.output_dim.""" | ||
| from torchembed.categorical import MultiCategoricalEmbedding | ||
|
|
Comment on lines
+150
to
+153
| 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) |
Comment on lines
+168
to
+171
| @property | ||
| def output_dim(self) -> int: | ||
| """Total output dimension of the categorical embeddings produced by this layer.""" | ||
| return self._cat_output_dim |
Comment on lines
+143
to
+146
| raise ImportError( | ||
| "The 'torchembed' package is required when embedding_backend='torchembed'. " | ||
| "Install it with: pip install torchembed" | ||
| ) from exc |
Comment on lines
+86
to
+87
| """torchembed backend must produce the same 2-D output shape as the native backend.""" | ||
| layer = Embedding1dLayer( |
Comment on lines
+115
to
+116
| """Gradients must flow through the torchembed embedding layer.""" | ||
| layer = Embedding1dLayer( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
embedding_backendfield toModelConfig(choices:"native"|"torchembed", default"native") so users can opt in totorchembed(v0.3.1+) as the categorical embedding engine with a single config change."torchembed"is selected,Embedding1dLayerdelegates totorchembed.categorical.MultiCategoricalEmbedding, which fuses all per-column embeddings into a single module with auto-sized dimensions.torchembedis an optional dependency (pip install pytorch-tabular[torchembed]); the native path is completely unchanged.Motivation
pytorch-tabular currently builds one
nn.Embeddingper categorical column and requires the user to either accept themin(50, (cardinality+1)//2)heuristic or supply explicitembedding_dims. torchembed exposes the same heuristic through a cleaner, single-module API (MultiCategoricalEmbedding) that:output_dimattribute instead of requiring callers to sum over a list of tuplesUsage
Files changed
src/pytorch_tabular/models/common/layers/embeddings.pyEmbedding1dLayer: newembedding_backendparam,output_dimproperty, torchembed forward pathsrc/pytorch_tabular/config/config.pyModelConfig: newembedding_backendfield + docstringsrc/pytorch_tabular/models/category_embedding/category_embedding_model.pyembedding_backendfrom configpyproject.toml[project.optional-dependencies] torchembedgrouptests/test_torchembed_backend.pyNotes
CategoryEmbeddingModelwires the option end-to-end in this PR. Other models that useEmbedding1dLayer(DANet,GANDALF,GATE,NODE) can be updated the same way in follow-up PRs — the layer-level API is already there.embedding_backend="torchembed"is set andtorchembedis not installed, a clearImportErrorwith an install command is raised atEmbedding1dLayer.__init__time, not silently at forward time.Test plan
pytest tests/test_torchembed_backend.py -v— 7/7 pass (torchembed installed)pytest.importorskip) when torchembed is absent, so existing CI remains green with no new required dependency🤖 Generated with Claude Code