Large refactor that solidifies object relationships#5
Merged
Conversation
- Aiming for 1:1 parity with umap_pytorch - Written with AI - Not tested
- Also remove callbacks.py (use default checkpointing)
- Merge UMAPConfig and TrainConfig directly into ParametricUMAP, removing nested config classes and the `create` factory method - Add device_to_lightning_acceleration_config utility for proper GPU device handling in Lightning Trainer - Remove unused set_global_seeds in favor of pl.seed_everything for reproducibility
- Simplify GlassBoxUMAP by extending ParametricUMAP dataclass instead of wrapping it - Rename DeepReLUNet to DeepPReLUNet with configurable n_hidden_layers parameter - Extract LayerNormDetached to new components.py module - Streamline compute_attributions method with cleaner batch Jacobian computation - Fix mutable default argument in ConvEncoder hidden_dims - Add tests for DeepPReLUNet architecture and device acceleration config
ekiefl
marked this pull request as ready for review
February 4, 2026 23:23
james-golden-arcadia
approved these changes
Feb 5, 2026
james-golden-arcadia
left a comment
james-golden-arcadia
left a comment
Collaborator
There was a problem hiding this comment.
Thank you for making all of these excellent changes! Runs smoothly for me. Repulsion_strength propagates correctly. The start on unit tests is very nice to have.
There are some small updates I would like to make in a next PR, mainly around running the visualize script for MNIST with a larger dataset and an optional number of works.
Overall LGTM.
Collaborator
Author
|
Interesting that windows 3.13 fails, but I'm gonna force merge anyways. |
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
2 participants
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
Major refactor of the parametric UMAP implementation with improved architecture, better code organization, and comprehensive test coverage.
GlassBoxUMAPto inherit fromParametricUMAP, reducing code duplicationDefaultEncoder,ConvEncoder)DeepPReLUNetwith dropout andLayerNormDetachedfor improved training stabilityParametricUMAPconfig into a single dataclass for cleaner API