Skip to content

fix(preprocessor): validate embeddings against the fitted dimensions #34

Description

@ChrisW09

Summary

Preprocessor.fit records embedding_dimensions_ for every embedding array it is given,
but transform never checks anything against it. Mismatched embeddings are accepted
silently, including ones with the wrong number of rows — which yields a result dict whose
blocks have different heights.

Reproduction

import numpy as np, pandas as pd
from pretab import Preprocessor

rng = np.random.default_rng(0)
df = pd.DataFrame({"a": rng.normal(size=100), "c": rng.choice(list("xyz"), 100)})
y = rng.normal(size=100)

pre = Preprocessor(numerical_method="minmax").fit(df, y, embeddings=rng.random((100, 8)))
print(pre.embedding_dimensions_)

print(pre.transform(df, embeddings=rng.random((100, 3)))["embedding_1"].shape)
print({k: v.shape for k, v in pre.transform(df, embeddings=rng.random((7, 8))).items()})
{'embedding_1': 8}
(100, 3)
{'num_a': (100, 1), 'cat_c': (100, 1), 'embedding_1': (7, 8)}

The third line is the damaging one: a caller stacking those blocks gets a shape error far
from the cause, or — worse — silently misaligned rows if they index rather than stack.

Two more unchecked cases:

pre2 = Preprocessor(numerical_method="minmax").fit(
    df, y, embeddings=[rng.random((100, 4)), rng.random((100, 5))])
print(pre2.embedding_dimensions_)
print(sorted(pre2.transform(df, embeddings=[rng.random((100, 4))])))
{'embedding_1': 4, 'embedding_2': 5}
['cat_c', 'embedding_1', 'num_a']      # embedding_2 silently missing

And fitting with embeddings then transforming without them drops the blocks with no
complaint, even though the reverse direction is correctly rejected with
IncompatibleParamsError.

Expected

transform validates each embedding against what fit recorded — count, per-array width,
and row count against X — and raises a pretab error naming the mismatch.

Actual

Nothing is validated. embedding_dimensions_ is written at fit and never read.

Root cause

pretab/preprocessor.py:472-483:

if embeddings is not None:
    if not self.embeddings_:
        raise IncompatibleParamsError(...)
    if isinstance(embeddings, np.ndarray):
        transformed_dict["embedding_1"] = embeddings.astype(np.float32)
    elif isinstance(embeddings, list):
        for idx, e in enumerate(embeddings):
            transformed_dict[f"embedding_{idx + 1}"] = e.astype(np.float32)

The only check is "were embeddings expected at all". fit
(pretab/preprocessor.py:359-367) populates embedding_dimensions_ purely as metadata.

Suggested fix

Normalize both fit and transform to a list, then validate in transform:

  • number of arrays matches len(self.embedding_dimensions_);
  • each array's shape[1] matches the recorded dimension;
  • each array's shape[0] matches X.shape[0];

raising PretabDataError naming the offending index and the expected vs actual shape. The
asymmetry around omitting embeddings entirely should also be settled — either raise (matching
the opposite direction) or document that they are optional at transform time.

Impact

Preprocessor.transform(..., embeddings=...). This is the documented integration point for
an embedding host such as DeepTab, which is exactly the caller most likely to wire up the
wrong array and least likely to notice.

Environment

  • pretab 0.1.0 (main @ 51c3043)
  • Python 3.11.15, numpy 2.4.6, pandas 2.3.3, scikit-learn 1.9.0, scipy 1.17.1
  • macOS (darwin 25.5.0)

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions