Size tensor embeddings without re-processing samples

[solarsys]

Problem

For TensorProcessor features, EmbeddingModel.__init__ inferred the input width by calling processor.process(sample[field]) on the first dataset sample (pyhealth/models/embedding.py). But samples from a SampleDataset are already processed, so process() ran a second time.

That's harmless for the plain TensorProcessor. It's wrong for any processor that isn't idempotent: one that selects columns, scales or imputes. Re-processing either raises, or silently produces the wrong width. On master, a column-selecting subclass raises inside MLP(dataset=...).

Changes

• EmbeddingModel reads the width from the first already-processed sample, which is exactly what the model will receive. It falls back to processor.size() when no sample is available. process() is never called.
• TensorProcessor.fit() records the feature width (the last dimension, or 1 for scalars), and size() returns it instead of None. Nothing in the library relied on None (checked by grep). Old pickled processors without the attribute still return None.

The processed sample takes precedence over size() on purpose. A subclass that changes the width but doesn't override size() would otherwise report the raw width.

Tests, docs, example

• New tests/core/test_embedding_tensor_width.py:

  • size() after fit(), before fit, and for scalars;
  • a column-selecting processor whose process() raises if re-run builds an MLP with in_features == 2;
  • plain "tensor" features keep their width.

  It fails on master and passes here.

• docs/api/processors/pyhealth.processors.TensorProcessor.rst gains a "Custom tensor processors" section. The EmbeddingModel and TensorProcessor docstrings get >>> examples, verified as doctests.

• New examples/custom_tensor_processor.py: a column-selecting processor with an MLP on synthetic data.

• Full core suite: Ran 1386 tests … OK (skipped=76). tools/check_pr_rules.py passes.

Reported by a downstream EHR project that uses a column-selecting, standardising tensor processor.

🤖 Generated with Claude Code