Skip to content

Make GrapPA feature outputs configurable#123

Merged
francois-drielsma merged 1 commit intomainfrom
feature/grappa-save
Apr 9, 2026
Merged

Make GrapPA feature outputs configurable#123
francois-drielsma merged 1 commit intomainfrom
feature/grappa-save

Conversation

@francois-drielsma
Copy link
Copy Markdown
Member

Description

This PR makes GrapPA feature outputs configurable instead of returning them unconditionally.

Specifically, it:

  • adds a return_features configuration flag to GrapPA
  • only includes node_features, edge_features, and global_features in the model output when that flag is enabled
  • keeps the internal forward-pass behavior unchanged, so the features are still computed and used by the GNN as before
  • includes a few small cleanup fixes in grappa.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

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation update
  • Code quality improvement (refactoring, type hints, etc.)
  • Performance improvement

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.

  • Verified that node_features, edge_features, and global_features are only added to the result when return_features=True
  • Verified that the GrapPA forward pass still computes and uses the feature tensors regardless of output configuration
  • Ran formatting/lint hooks during commit (black, isort, flake8, pre-commit checks)

Test Configuration:

  • OS: macOS
  • Python version: local development environment
  • SPINE version: feature/grappa-save, commit 94820c67

To reproduce:

  1. Run a GrapPA-based configuration with return_features: false and confirm that the output dictionary does not include node_features, edge_features, or global_features.
  2. Run the same configuration with return_features: true and confirm that those feature tensors are present in the result.
  3. If using full-chain mode, confirm that the namespaced GrapPA outputs remain available as expected when the flag is enabled.

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings or errors
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • Any dependent changes have been merged and published

Screenshots / Event Displays (if applicable)

None included.

Additional Notes

A few implementation details are worth calling out for reviewers:

  • This PR changes output policy, not model semantics: the feature tensors are still computed and fed into the GNN exactly as before.
  • The main effect is to avoid returning large intermediate feature tensors by default.
  • This is particularly useful in full-chain inference, where returning intermediate GrapPA features by default can significantly increase output payload size.
  • The flag lives directly on the GrapPA model configuration, so feature dumping remains available when explicitly requested.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FR] Add intermediate "feature-level" output for validation workflows and detector systematic studies

1 participant