infer/predict mis-assign every train-split output: run_predict_step pairs schema_df order with a shuffled dataloader

[florian6973]

_runtime.run_predict_step pairs the index with model outputs on the assumption that the dataloader
iterates in schema_df order:

# commands/_runtime.py
#: MEDS split name → the meds-torch-data `Datamodule` (dataset, dataloader) attribute pair. Only the
#: non-shuffling loaders are listed: prediction order must match `schema_df` order.
SPLIT_ATTRS = {"train": ("train_dataset", "train_dataloader"), ...}
...
keys = dataset.schema_df.select(KEYS)
return keys, trainer.predict(module, dataloaders=dataloader)

That holds for tuning and held_out. It does not hold for train:

# meds_torchdata/extensions/lightning_datamodule.py
def train_dataloader(self):
    return self.__dataloader(self.train_dataset, shuffle=True)

meds-torch-data says so in its own class docstring — "The train dataloader shuffles so doesn't return
stable outputs, but the others do not shuffle." Lightning notices (Your predict_dataloader's sampler has shuffling enabled…) but only warns.

So every train-split row of <data_dir>/inference/<name>/artifacts.parquet is written against one
subject's key while carrying another subject's embedding, permuted differently on every run.

Why it is invisible

The artifact has the right keys, the right row count, the right vector width, a valid
kind: embeddings manifest, and it passes the probe's coverage check. test_end_to_end scores
splits=[held_out], which is one of the two splits that happen to be correct.

Measured

Reference implementation: a probe-profile port of PORTER (arXiv:2606.24102). Fixture: the template's
own build_signal_dataset (200 train / 40 tuning / 40 held-out subjects, label = presence of
SIGNAL//POS). One fixed pretrained checkpoint; 5-fold CV AUROC of a standardised logistic regression
fitted on the materialized embeddings of each split:

split	as shipped	with a non-shuffling loader
train	0.49 (chance)	1.00
tuning	1.00	1.00
held_out	1.00	1.00

End to end: a probe fitted on the scrambled train split scored AUROC 0.461 on held-out; one fitted on
the correctly-aligned tuning split scored 1.000, from the same embeddings and the same frozen
backbone.

How it surfaced

Not from a failing test — pytest -m slow passed. Two runs at the same seed reported different signal
AUROCs (0.9427, then 1.0000). Bisecting: preprocess_data reproduces byte-for-byte, pretrain
bit-for-bit (identical val/loss to 14 digits, identical weight digests at 6 and 25 epochs),
supervised_train and predict bit-for-bit from fixed inputs — and infer did not, on exactly 200 of
280 rows. 280 − 200 = 80 = the tuning and held-out splits.

Worth stating: the learnability tier had been passing for the wrong reason. The probe trains on
train and validates on tuning, so ModelCheckpoint(monitor=val/loss) selected whichever epoch scored
best against the clean validation split — which is why it still reached 0.94–1.00, and why it landed
somewhere different each run.

Scope

• Every probe-profile model, identically.
• SupervisedPredictCommand shares run_predict_step, so any model predicting with splits=[train]
  (or the default "every split present") is affected too.

Suggested fix

Build the loader from the dataset rather than reading train_dataloader() off the datamodule:

dataset = getattr(datamodule, SPLIT_ATTRS[split][0])
loader = dataset.get_dataloader(batch_size=..., num_workers=..., shuffle=False)

and add an alignment assertion to the conformance suite — test_end_to_end predicting splits=[train],
or a model-free check that batch row i holds the timeline of the subject schema_df lists at row i.
The latter is what the port uses; it needs no model and catches this directly.

Workaround in the meantime: subclass DefaultInferCommand in the model's own commands.py.