[dalloliogm]

Hello,

Almost one year ago I started contributing to the VAE bounty in PyHealth. Unfortunately, I messed up with Git and became busy with work, and time passed, so I never followed up on the review comments on my previous PR.

This is a new implementation of the bounty for PyHealth 2.0.

Contributor: Giovanni Dall’Olio

Contribution Type: Model

Description: The VAE now supports image inputs and categorical EHR code sequences. I’ve also added or updated three notebooks demonstrating synthetic image generation, ChestXray14 image modeling, and EHR sequence modeling, with explanations of potential applications.

Files to Review:

• pyhealth/models/vae.py - VAE implementation for image and categorical sequence inputs.
• tests/core/test_vae.py - model behavior, generation, and Trainer compatibility tests.
• examples/cxr/vae_cxr_synthetic.ipynb - synthetic image VAE example.
• examples/cxr/vae_cxr_chestxray14.ipynb - ChestXray14 image example.
• examples/mortality_prediction/vae_ehr_sequences.ipynb - EHR sequence example.
• examples/cxr/chextXray_image_generation_VAE.py - updated image VAE example.
• docs/tutorials.rst and docs/api/models/pyhealth.models.VAE.rst - tutorial links and API documentation.

[github-actions]

Test Coverage

Total: 75.19% (21552 / 28664 lines)

Modified file	Coverage	Lines
pyhealth/calib/predictionset/favmac/__init__.py	24.32%	27 / 111
pyhealth/calib/predictionset/favmac/core.py	18.58%	21 / 113
pyhealth/datasets/sample_dataset.py	94.48%	325 / 344
pyhealth/datasets/utils.py	45.31%	87 / 192
pyhealth/medcode/cross_map.py	96.91%	94 / 97
pyhealth/models/__init__.py	100.00%	54 / 54
pyhealth/models/base_model.py	94.33%	133 / 141
pyhealth/models/gamenet.py	89.45%	178 / 199
pyhealth/models/rnn.py	93.77%	346 / 369
pyhealth/models/sequence_embedding.py	100.00%	148 / 148
pyhealth/models/vae.py	85.96%	202 / 235
pyhealth/processors/__init__.py	94.29%	33 / 35
pyhealth/processors/base_processor.py	79.73%	59 / 74
pyhealth/processors/sequence_processor_v2.py	100.00%	683 / 683
pyhealth/processors/timeseries_processor.py	88.35%	91 / 103
pyhealth/tasks/chestxray14_binary_classification.py	100.00%	22 / 22
pyhealth/tasks/yaib_mimic4.py	93.06%	523 / 562
pyhealth/trainer.py	76.38%	207 / 271

Workflow run · f2973cd

[jhnwu3]

/review

[dalloliogm]

If the PR is too large, I can split it into two, so the AI review can take place.

The changes to the VAE class are small. Most of the new lines come from the notebooks included to describe what is a VAE and how it can be used. I could split those into a different PR

[jhnwu3]

Yes, please split it if possible.

[jhnwu3]

I think the main thing is the .ipynb notebooks are massive, so just one with a simple examples *.py file and then the follow up notebooks would be much appreciated.

[dalloliogm]

Hi, I think the problem was that I forgot to clean up the notebook outputs. I've done that now - can you try reviewing again?

There were a few plots that made the notebooks very big - but they should not have been committed.

[dalloliogm]

/review

[github-actions]

Previous review

_{Reviewed commit 3e0ff51}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: human review

Thank you for coming back to the VAE bounty and porting it to PyHealth 2.0. Adding categorical EHR sequence support, optional label conditioning for images, and three example notebooks is a useful extension. Line coverage for pyhealth/models/vae.py is 88.29%, above the 80% bar, and total repository coverage is 74.07%, above the 60% bar.

A maintainer will also take a look. The change alters the public contract of the existing VAE: hidden_dim changes meaning and default, latent_dim is new, and the generate signature changes. Existing correct calls can therefore break or behave differently, and a maintainer should decide whether that is acceptable.

To get this merged:

1. Single-label image datasets. The constructor auto-uses a sole output label as the condition. A dataset with a single multilabel or regression output, for example ChestXray14 multilabel, now raises an error, where the old version ignored labels. Only auto-condition when the label is binary or multiclass, and otherwise ignore it unless condition_key is given explicitly.
2. generate(num_samples=...) on labeled image datasets. This now raises because of the auto-conditioning. Please keep unconditional generation working, for example by sampling the condition at random or only conditioning when condition_key is passed explicitly.
3. @​torch.no_grad() removed from generate. The returned images now require grad, so .numpy() and similar calls fail. Please restore the decorator or detach the output. The old device argument has also been dropped, and the second positional argument is now sequence_length. Please keep the signature backward compatible or call the change out clearly.
4. hidden_dim semantics. Previously hidden_dim was the latent size, with a default of 128. It is now the channel width, with a default of 32. Existing VAE(dataset, hidden_dim=N) calls silently build a different network. Please keep hidden_dim as the latent size, or document the break prominently in the description and docs.
5. Tests. The COVID19CXRImageGeneration task tests and the learning test were deleted, although that task still ships. Please restore tests for the task and a test that training reduces the reconstruction error. Also add tests for the new error branches and for Trainer.train.

Optional suggestions:

• In the ChestXray14 notebook, the comment about patient-level splits being streaming subsets looks inaccurate. split_by_patient returns subsets that get_dataloader should accept, so you could use it directly.
• Unsupported PIL modes such as "P" raise a bare KeyError at the mode lookup in _init_image. A clear ValueError would be friendlier.

_{Reviewed commit 3e0ff51}

[dalloliogm]

/review

[dalloliogm]

/review

[github-actions]

Previous review

_{Reviewed commit 478feab}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. It extends the PyHealth 2.0 VAE to categorical EHR code sequences and optional label-conditioned image generation, and it adds three example notebooks and docs updates. The new model code looks reasonable. A maintainer will also take a look, because the rewrite of pyhealth/models/vae.py changes the existing VAE's public interface and the format of saved checkpoints.

Coverage: pyhealth/models/vae.py is at 89.05%, above the 80% bar for modified files, and the repository total is 74.11%, above the 60% bar. I did not run the tests or notebooks.

To get this merged:

1. Keep backward compatibility, or state the break clearly. The old encode() and decode() methods are gone. The attributes encoder, decoder, mu and logvar are renamed, and the ResBlock2D architecture is replaced. State_dicts saved from the previous VAE therefore no longer load. Please restore encode and decode (or add thin wrappers). If the architecture change is intentional, document it in the PR description and the docs as a breaking change.
2. Image processors without an explicit mode used to work, because the channel count was read from the data. They now raise a ValueError. Please infer the channel count from a sample when mode is None, or document this as an intentional break. Also keep a clear error for to_tensor=False, which the old to_tensor check used to catch.
3. Restore the test coverage that was removed. The multi-epoch "learns to reconstruct" test is now a single optimizer step. The RGB, gradient-flow and state_dict round-trip tests are gone, so the RGB path used by the COVID classification example is untested. Please add at least an RGB forward-pass test and a state_dict round-trip test.

Optional suggestions:

• The sequence path always samples the latent, even in eval mode, so validation loss is noisy. Using mu when not self.training would make it deterministic.
• In the synthetic notebook, torch.manual_seed before an eval-mode forward has no effect, so the remark about seeding can be dropped.
• The ChestXray14 notebook subclasses the binary task only to add patient_id. That would be better placed in the task itself.

_{Reviewed commit 478feab}

[dalloliogm]

I just noticed that the image-generation part of this bounty was already covered by PR #1351. Thanks for getting that implementation merged.

This PR builds on that work by extending the VAE to categorical EHR sequences, alongside image inputs. It adds examples for synthetic image generation, conditional image generation with ChestXray14, and EHR sequence modeling, and expands the tests to cover reconstruction, generation, and Trainer compatibility across both input types.

I’ve kept image conditioning opt-in and preserved the existing image VAE behavior and COVID-19 CXR task from #1351. The current CI checks are passing. My aim is to make this a complementary extension to #1351, rather than duplicate its image-generation work.

[dalloliogm]

/review

[jhnwu3]

/review

[github-actions]

_{Reviewed commit f2973cd}

[github-actions]

[medium] Image mode whitelist rejects modes the previous VAE accepted

The new check raises for any ImageProcessor mode outside L/RGB/RGBA. The previous VAE had no mode check: channels are read from the processed sample's shape, so modes such as LA (2 channels) or CMYK (4 channels) worked and now fail at construction. Nothing in the encoder or decoder requires the restriction, and the docs do not mention it. Fix: remove the mode check and rely on the sample-shape check, or document the restriction and keep it only if it is necessary.

[github-actions]

[low] Sequence VAE output lacks y_true/y_prob, so Trainer.inference fails

_forward_sequence returns loss, reconstruction_logits, mu, log_var and mask, but no y_true or y_prob. Trainer.inference (pyhealth/trainer.py around line 391) reads output['y_true'] and output['y_prob'], so it raises KeyError for sequence VAEs. Trainer.train and Trainer.evaluate work. Fix: return y_true and y_prob (for example the tokens and softmax probabilities), or state in the docstring that inference is unsupported.

[github-actions]

[low] Ambiguous-feature error and minimum image size check are no longer tested

The earlier tests for the multiple-input error and the images-smaller-than-8-pixels error were removed, and the new test file has no replacement, so these new and existing validation paths have no coverage. Fix: add a test that a dataset with two supported inputs and no feature_key raises ValueError, and a test that a 4x4 image raises ValueError.

[github-actions]

[low] self.image_size is assigned twice and never read

The attribute is set from the processor, then overwritten with height in the height != width branch, and nothing else in the class reads it. Remove both assignments, or use it. The branch comment describes behavior the code does not need.

[github-actions]

[low] Sequence VAE builds embedding tables for every input feature

EmbeddingModel(dataset, ...) creates a table for each sequence input in the dataset, but _forward_sequence uses only feature_key. For multi-feature tasks, the unused tables are included in model.parameters() and in the optimizer. Fix: create a single nn.Embedding for the chosen processor, or document the extra parameters.

[github-actions]

[low] Example runs at import time with no evaluation or main guard

The script trains at module level and prints generated shapes, with no train/validation split or reconstruction check, and the previous example's comparison plots were removed. Wrap the code in main() under an if name == 'main' guard, and consider keeping a reconstruction check so the example demonstrates the model's quality.

[github-actions]

AI Review

suggestion: need change
decision: need change

Thank you for coming back with a fresh PyHealth 2.0 implementation. This PR extends the VAE with a categorical code-sequence path and optional label conditioning for images. It also adds three example notebooks, updated docs, and a patient_id key in ChestXray14BinaryClassification so patient-level splits work. Keeping the image encoder/decoder/mu/logvar module names, the encode()/decode() methods and the hidden_dim latent size is a good choice for compatibility. Coverage meets the bar: pyhealth/models/vae.py is at 85.96% (bar is 80%) and the whole repository is at 75.19% (bar is 60%).

To get this merged:

• Remove the new image-mode restriction in pyhealth/models/vae.py (around line 172). The check rejects any ImageProcessor mode other than "L", "RGB" and "RGBA". The previous VAE took the channel count from a processed sample and accepted any mode, such as "LA" or "CMYK". Configurations that worked before would now fail, and I can't see a technical reason for it. Please drop the check, or widen it to match what the model already handled.
• Update test_unsupported_image_mode_raises_value_error to match. It currently asserts the restriction.

Optional suggestions:

• The sequence VAE doesn't return y_true or y_prob, so Trainer.inference() raises a KeyError on it. A short docstring note, or a graceful path, would help.
• Two old tests no longer exist: the "multiple inputs without feature_key" error and the 8-pixel minimum-size check. Please restore them.
• self.image_size is assigned twice and never used in vae.py, so it can go.
• EmbeddingModel allocates tables for every input feature, not only the sequence one.
• examples/cxr/chextXray_image_generation_VAE.py now runs at import time, with no main() guard and no evaluation step. A guard would make it friendlier to import and test.

_{Reviewed commit f2973cd}