Fix(geotransolver): fail fast when include_local_features=True without geometry#1766
Fix(geotransolver): fail fast when include_local_features=True without geometry#1766ShubhJain007 wants to merge 1 commit into
Conversation
…t geometry
GeoTransolver sizes its GALE blocks for
effective_hidden = n_hidden + n_hidden_local * len(radii) whenever
include_local_features=True, but the local multi-scale features are only
produced when a geometry tensor is supplied. Omitting geometry/geometry_dim
silently skipped the concat and surfaced as a confusing downstream
LayerNorm shape mismatch (e.g. normalized_shape=[192] vs input 128).
Add two fail-fast guards with actionable messages:
- GeoTransolver.__init__: include_local_features=True requires geometry_dim.
- GlobalContextBuilder.build_context: geometry is required at forward time
when local features are enabled.
Add a regression test covering both guards and the working path.
Signed-off-by: Shubh Jain <shubhjain10102003@gmail.com>
Greptile SummaryThis PR adds two fail-fast input-validation guards to GeoTransolver to surface a previously silent misconfiguration — using
Important Files Changed
Reviews (1): Last reviewed commit: "Fix(geotransolver): fail fast when inclu..." | Re-trigger Greptile |
| # Extract multi-scale features if enabled | ||
| if self.local_extractors is not None and geometry is not None: |
There was a problem hiding this comment.
The
and geometry is not None guard in this condition is now logically redundant. The new guard added above already raises ValueError whenever self.local_extractors is not None and geometry is None, so by the time execution reaches this line, geometry is guaranteed to be non-None whenever local_extractors is set. Simplifying the condition removes the dead branch and makes the intent clearer.
| # Extract multi-scale features if enabled | |
| if self.local_extractors is not None and geometry is not None: | |
| # Extract multi-scale features if enabled | |
| if self.local_extractors is not None: |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| assert not torch.isnan(compiled_out).any() | ||
|
|
||
|
|
||
| def test_geotransolver_local_features_requires_geometry(device): |
There was a problem hiding this comment.
The similar test
test_geotransolver_local_features_compile directly above is decorated with @requires_module("warp"). The new test omits that decorator, which is fine for cases (A) and (B) since they raise before any BQWarp call, and radius_search does have a pure-PyTorch fallback (rank=1, baseline). However, if CI enforces the same guard for all local-features tests for consistency, adding the decorator (or a comment explaining the intentional omission) would prevent future confusion.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
PR: GeoTransolver — fail fast when
include_local_features=Truewithout geometryRepo: NVIDIA/physicsnemo
Area:
physicsnemo/experimental/models/geotransolver/Related: issue #1606 (also an input-validation gap on the GeoTransolver local path; this PR is a distinct fix — missing-geometry guard, not the ball_query positions-dim check).
Problem
With
include_local_features=True, the GALE blocks are sized foreffective_hidden = n_hidden + n_hidden_local * len(radii), but the extra(local multi-scale) channels are only produced when a
geometrytensor issupplied. If
geometry_dim/geometryis omitted, the concat is silently skippedand the token stream stays at
n_hidden, producing a confusing downstream error:Fix
Fail fast with an actionable message, in two places:
GeoTransolver.__init__):include_local_features=Truerequiresgeometry_dimto be set.GlobalContextBuilder.build_context: if local extractors are active,geometrymust be passed at forward time.
( 2 small guards, no behavior change for valid usage.)