Make GrapPA feature outputs configurable#123
Merged
francois-drielsma merged 1 commit intomainfrom Apr 9, 2026
Merged
Conversation
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
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.
Description
This PR makes GrapPA feature outputs configurable instead of returning them unconditionally.
Specifically, it:
return_featuresconfiguration flag toGrapPAnode_features,edge_features, andglobal_featuresin the model output when that flag is enabledgrappa.py(typos/docstring consistency)Motivation and Context
GrapPA feature tensors can be useful for debugging, analysis, and feature dumping, but they are relatively heavy outputs and are not needed in the default workflow. Returning them unconditionally increases output size, memory pressure, and numpy conversion/serialization overhead, especially in full-chain inference where GrapPA outputs are namespaced and merged into the global result dictionary.
This change makes feature dumping explicit and opt-in, while preserving the ability to expose those features when needed. In particular, it addresses the concern raised in issue #117 about exposing GrapPA-produced features without forcing them into every default output.
Fixes #117
Type of Change
How Has This Been Tested?
I reviewed the GrapPA forward path and verified that the feature tensors are still used internally by the GNN, while only their exposure in the returned result dictionary is gated by configuration.
node_features,edge_features, andglobal_featuresare only added to the result whenreturn_features=Trueblack,isort,flake8, pre-commit checks)Test Configuration:
feature/grappa-save, commit94820c67To reproduce:
return_features: falseand confirm that the output dictionary does not includenode_features,edge_features, orglobal_features.return_features: trueand confirm that those feature tensors are present in the result.Checklist
Screenshots / Event Displays (if applicable)
None included.
Additional Notes
A few implementation details are worth calling out for reviewers:
GrapPAmodel configuration, so feature dumping remains available when explicitly requested.