Mainline geo t by coreyjadams · Pull Request #1845 · NVIDIA/physicsnemo

coreyjadams · 2026-07-16T21:57:19Z

PhysicsNeMo Pull Request

This PR moves GeoTransolver out of experimental and into physicsnemo proper. Because GeoT also has a FLARE attention backend, we also move FLARE at the same time.

A few other pieces change with this:

FLARE now has full TransformerEngine support.
GeoTransolver's OOD component is transformed from an integrated-directly piece of the model to a wraps-the-model component.
GALE, flare_attention move to physicsnemo.nn
Context_projector stays with GeoTransolver in the models folder.

Tests are updated, and most examples are changed only to update import paths, etc. The major update is the OODGuard changes. No example yet is exercising the FLARE TE path.

A docs update will come in a follow on PR so it can be reviewed on a separate time scale.

Description

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.
The CHANGELOG.md is up to date with these changes.
An issue is linked to this pull request.
If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

…are not yet updated.

copy-pr-bot · 2026-07-16T21:58:43Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-07-16T22:06:22Z

Greptile Summary

This PR promotes GeoTransolver and FLARE from physicsnemo/experimental/models to production (physicsnemo/models), moves GALE, GALE_FA, and FLARE attention layers into physicsnemo/nn/module, and refactors the embedded OOD guard into a standalone GuardedGeoTransolver wrapper in physicsnemo/experimental/guardrails.

GeoTransolver / GALE promotion: The model and its attention blocks are moved to production with accompanying non-regression tests, checkpoint tests, and reference .pth files. The guard_config parameter is removed (acceptable since it was still in experimental) and replaced with the external GuardedGeoTransolver wrapper. A new GALE_FA attention variant using FLARE self-attention is added.
FLARE promotion: FLARE model and attention layer are moved to production, registered in pyproject.toml, and covered by a new comprehensive test suite.
Bug in MultiScaleFeatureExtractor.extract_local_features: In context_projector.py, processor(geometry, spatial_coords) swaps the argument order relative to the GeometricFeatureProcessor.forward(query_points, key_features) contract and relative to extract_context_features. When include_local_features=True, the ball-query search uses geometry feature values as spatial coordinates and aggregates positions instead of geometry — producing incorrect local features fed into the transformer.

Important Files Changed

Filename	Overview
physicsnemo/models/geotransolver/context_projector.py	New file (promoted from experimental). Contains a P1 argument-order bug in `MultiScaleFeatureExtractor.extract_local_features` — `processor(geometry, spatial_coords)` is swapped relative to `extract_context_features` and the `GeometricFeatureProcessor.forward(query_points, key_features)` contract, producing incorrect results when `include_local_features=True`.
physicsnemo/models/geotransolver/geotransolver.py	Promoted from experimental to production. `concrete_dropout` parameter added with default but not documented in the class docstring Parameters section. Structure and validation logic look correct.
physicsnemo/nn/module/gale.py	Moved from experimental to physicsnemo/nn/module. Adds GALE_FA and _GALEStructuredForwardMixin. `GALE_FA.__init__` is missing a type annotation for the `dim` parameter (MOD-006).
physicsnemo/experimental/guardrails/embedded/geotransolver_guard.py	New experimental wrapper. Forward-hook approach for capturing geometry latent is sound. Class docstring uses `"""` instead of `r"""` (MOD-003b) and inherits `nn.Module` instead of `physicsnemo.Module` (MOD-001).
physicsnemo/models/flare/flare.py	Promoted from experimental. FLARE model correctly inherits from Transolver and replaces physics-attention blocks with _FLAREBlock. Logic and structure are correct.
physicsnemo/nn/module/flare_attention.py	Moved to physicsnemo/nn/module. Two-pass FLARE self-attention is correctly implemented. Forward/Outputs docstring sections use non-LaTeX shape notation (MOD-003e).
test/models/geotransolver/test_geotransolver.py	`@requires_module("warp")` removed from local-features tests; correct since `BQWarp`/`radius_search` has a pure-torch fallback. Embedded guard tests moved to new file. New TE + GALE_FA test added.
test/experimental/guardrails/embedded/test_geotransolver_guard.py	New test file for GuardedGeoTransolver. Covers collect/check cycle, hook cleanup via close(), invalid configurations, and geometry-only / global-only surface combinations.
test/models/flare/test_flare.py	New test file moved from experimental. Covers constructor, 2D/irregular forward accuracy, CUDA graphs, JIT, AMP, checkpoint and ONNX deploy. Satisfies MOD-008a/b/c requirements.
pyproject.toml	Registers `FLARE` and `GeoTransolver` in the entry-points table for the physicsnemo model registry. No issues.

Comments Outside Diff (3)

physicsnemo/models/geotransolver/context_projector.py, line 755-761 (link)

Swapped arguments produce wrong local features

processor(geometry, spatial_coords) passes geometry as query_points (the spatial coordinates fed into BQWarp's neighbor search) and spatial_coords as key_features (the values aggregated by the MLP). This is the reverse of what extract_context_features does (processor(spatial_coords, geometry)), and the reverse of what GeometricFeatureProcessor.forward(query_points, key_features) expects. Whenever include_local_features=True, the concatenated local features fed into the transformer will be semantically incorrect.
physicsnemo/nn/module/gale.py, line 581-592 (link)

The dim parameter of GALE_FA.__init__ is missing a type annotation. Per coding standard MOD-006, all public method arguments must use type hints. All other parameters in this signature are annotated.

File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

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!
physicsnemo/models/geotransolver/geotransolver.py, line 387-411 (link)

concrete_dropout parameter undocumented in class docstring (MOD-003a)

The concrete_dropout: bool = False parameter is present in __init__ but is not listed in the Parameters section of the class docstring. Per coding standard MOD-003a/MOD-003d, every parameter must be documented.

File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

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!

_{Reviews (1): Last reviewed commit: "Merge branch 'main' into mainline-GeoT" | Re-trigger Greptile}

greptile-apps · 2026-07-16T22:06:26Z

+def _extract_global_embedding(args: tuple, kwargs: dict) -> torch.Tensor | None:
+    """Pull ``global_embedding`` from a GeoTransolver forward call.
+
+    ``GeoTransolver.forward`` takes ``global_embedding`` as its third positional


GuardedGeoTransolver should inherit from physicsnemo.Module (MOD-001)

GuardedGeoTransolver is a torch.nn.Module subclass. Per coding standard MOD-001, all model classes in this repository must inherit from physicsnemo.Module to benefit from serialisation, versioning, and registry features. Even as an experimental wrapper, this class exposes a forward() method and acts as a model.

File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

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!

greptile-apps · 2026-07-16T22:06:27Z

+class GuardedGeoTransolver(nn.Module):
+    """GeoTransolver wrapped with an out-of-distribution guard.


The GuardedGeoTransolver class docstring uses plain """ instead of the required raw-string prefix r""". Per coding standard MOD-003b, all class and method docstrings must use r""" so that LaTeX backslash sequences are not interpreted as Python escape sequences.

Suggested change

class GuardedGeoTransolver(nn.Module):

"""GeoTransolver wrapped with an out-of-distribution guard.

class GuardedGeoTransolver(nn.Module):

r"""GeoTransolver wrapped with an out-of-distribution guard.

File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

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!

peterdsharpe

Verified changes on the unified recipe (CODEOWNER), LGTM there!

One question - where is the new geotransolver_guard coming from? Is that splitting out functionality that used to be in core geotransolver.py?

CharlelieLrt · 2026-07-16T22:07:30Z

This PR is breaking the kinetic Monte-Carlo recipe as well as checkpoints that were distributed. Nothing against the move, but can we keep it importable from experimental at least for some time (maybe with warning or something)?

coreyjadams · 2026-07-16T22:07:47Z

One question - where is the new geotransolver_guard coming from? Is that splitting out functionality that used to be in core geotransolver.py?

Yes, exactly. This is really something that @mnabian needs to weigh in on, if we should build it this way or another way.

CharlelieLrt · 2026-07-16T22:11:51Z

Also, model hould be added to API doc I think?

coreyjadams · 2026-07-16T22:13:03Z

Also, model hould be added to API doc I think?

I would like to put in a second docs PR to give @megnvidia breathing room to review independent of code changes. I've got the updates staged locally already. Are you ok with that? We can have it ready up here before approval of this one.

coreyjadams · 2026-07-16T22:13:26Z

/ok to test 8a20c73

CharlelieLrt · 2026-07-16T22:14:45Z

-from .gale import (
-    GALE,
-    GALE_FA,
-    GALE_block,


nit: GALEBlock

coreyjadams · 2026-07-16T23:42:37Z

Docs update in #1847 . It's this whole PR + docs, so this one needs to merge first.

coreyjadams · 2026-07-17T13:34:27Z

@CharlelieLrt I've added shims back to experimental. @ktangsali maybe a good test is to try this PR in physicsnemo-cfd and validate the checkpoints import correctly? I can do it if you think that will be robust testing.

CharlelieLrt

Kinetic Monte-Carlo recipe to update with right import

megnvidia · 2026-07-17T18:18:41Z

 > By default, the normalization sets the mean to 0.0 and std to 1.0 of all labels in the dataset, computing the mean across the train dataset. You could adapt this to a different normalization, however take care to update both the preprocessing as well as inference scripts. Min/Max is another popular strategy.

-To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, see the API for `physicsnemo.models.transolver` and `physicsnemo.experimental.models.geotransolver`.
+To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, see the API for `physicsnemo.models.transolver` and `physicsnemo.models.geotransolver`.


Suggested change

To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, see the API for `physicsnemo.models.transolver` and `physicsnemo.models.geotransolver`.

To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, refer to the API for `physicsnemo.models.transolver` and `physicsnemo.models.geotransolver`.

megnvidia · 2026-07-17T18:19:33Z

-(leave `null` to disable):
+To enable the guard, wrap a constructed `GeoTransolver` with
+`GuardedGeoTransolver` (from `physicsnemo.experimental.guardrails.embedded`).
+The wrapper observes the two surfaces above via a forward hook and delegates the


Suggested change

The wrapper observes the two surfaces above via a forward hook and delegates the

The wrapper observes the two surfaces above through a forward hook and delegates the

megnvidia · 2026-07-17T19:02:57Z

+  `physicsnemo.models.geotransolver.GeoTransolver`, together with the FLARE
+  model (`physicsnemo.models.flare.FLARE`) and the reusable GALE and FLARE
+  attention layers (`physicsnemo.nn.GALE`, `physicsnemo.nn.GALE_block`,
+  `physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:


Suggested change

`physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:

`physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model.

megnvidia · 2026-07-17T19:03:26Z

+  model (`physicsnemo.models.flare.FLARE`) and the reusable GALE and FLARE
+  attention layers (`physicsnemo.nn.GALE`, `physicsnemo.nn.GALE_block`,
+  `physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:
+  wrap a constructed GeoTransolver with


Suggested change

wrap a constructed GeoTransolver with

Wrap a constructed GeoTransolver with

megnvidia · 2026-07-17T19:04:08Z

+  attention layers (`physicsnemo.nn.GALE`, `physicsnemo.nn.GALE_block`,
+  `physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:
+  wrap a constructed GeoTransolver with
+  `physicsnemo.experimental.guardrails.embedded.GuardedGeoTransolver` (or


Suggested change

`physicsnemo.experimental.guardrails.embedded.GuardedGeoTransolver` (or

GuardedGeoTransolver (or call

megnvidia · 2026-07-17T19:04:58Z

+  `physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:
+  wrap a constructed GeoTransolver with
+  `physicsnemo.experimental.guardrails.embedded.GuardedGeoTransolver` (or
+  `attach_ood_guard`) to enable out-of-distribution guarding — the


Suggested change

`attach_ood_guard`) to enable out-of-distribution guarding — the

`attach_ood_guard`) to enable out-of-distribution guarding. The `guard_config`

megnvidia · 2026-07-17T19:05:11Z

+  wrap a constructed GeoTransolver with
+  `physicsnemo.experimental.guardrails.embedded.GuardedGeoTransolver` (or
+  `attach_ood_guard`) to enable out-of-distribution guarding — the
+  `guard_config` model argument is removed.


Suggested change

`guard_config` model argument is removed.

model argument is removed.

mnabian

LGTM! Left a couple of minor comments. Great to see TE is now enabled for FLARE! Guardrail changes look good, but an action item for me is to update the user documentation accordingly.

mnabian · 2026-07-17T22:35:25Z

@@ -14,11 +14,10 @@
 # See the License for the specific language governing permissions and


You can also include a test for the return_embedding_states=True path.

mnabian · 2026-07-17T22:53:23Z

            self.attn_fn = te.DotProductAttention(
                num_attention_heads=self.heads,
                kv_channels=self.dim_head,
                attention_dropout=dropout,


When use_te=True, attention-weight dropout is applied inside the TE kernel in addition to the shared out_dropout, but the non-TE path skips attention dropout entirely. This won't affect our GeoTransolver trainings as we usually don't use dropout, but will affect the models that use Concrete Dropout for UQ, and makes the two modes non-equivalent for UQ.

coreyjadams added 6 commits July 16, 2026 12:49

Move GeoT and Flare and related layers out of experimental. Examples …

e3ac883

…are not yet updated.

Update all import paths in examples.

90cc4a7

Update entry points in pyproject.toml

b567c9a

Update the changelog

0baabcb

Enable TE through FLARE entirely

deadee3

Merge branch 'main' into mainline-GeoT

8a20c73

coreyjadams marked this pull request as ready for review July 16, 2026 21:58

coreyjadams requested review from RishikeshRanade, ktangsali, loliverhennigh, mnabian and peterdsharpe as code owners July 16, 2026 21:58

greptile-apps Bot reviewed Jul 16, 2026

View reviewed changes

peterdsharpe approved these changes Jul 16, 2026

View reviewed changes

CharlelieLrt reviewed Jul 16, 2026

View reviewed changes

coreyjadams mentioned this pull request Jul 16, 2026

GeoT docs update #1847

Draft

6 tasks

Add shims to not break old checkpoints

6f13af0

CharlelieLrt requested changes Jul 17, 2026

View reviewed changes

megnvidia approved these changes Jul 17, 2026

View reviewed changes

mnabian approved these changes Jul 17, 2026

View reviewed changes

		class GuardedGeoTransolver(nn.Module):
		"""GeoTransolver wrapped with an out-of-distribution guard.

	To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, see the API for `physicsnemo.models.transolver` and `physicsnemo.models.geotransolver`.
	To configure your training run, use `hydra`. The config contains sections for the model, data, optimizer, and training settings. For details on the model parameters, refer to the API for `physicsnemo.models.transolver` and `physicsnemo.models.geotransolver`.

	The wrapper observes the two surfaces above via a forward hook and delegates the
	The wrapper observes the two surfaces above through a forward hook and delegates the

	`physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model:
	`physicsnemo.nn.FLARE`). The embedded OOD guard is decoupled from the model.

	wrap a constructed GeoTransolver with
	Wrap a constructed GeoTransolver with

	`physicsnemo.experimental.guardrails.embedded.GuardedGeoTransolver` (or
	GuardedGeoTransolver (or call

	`attach_ood_guard`) to enable out-of-distribution guarding — the
	`attach_ood_guard`) to enable out-of-distribution guarding. The `guard_config`

	`guard_config` model argument is removed.
	model argument is removed.

		@@ -14,11 +14,10 @@
		# See the License for the specific language governing permissions and

Uh oh!

Conversation

coreyjadams commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PhysicsNeMo Pull Request

Description

Checklist

Dependencies

Review Process

Uh oh!

copy-pr-bot Bot commented Jul 16, 2026

Uh oh!

greptile-apps Bot commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Important Files Changed

Comments Outside Diff (3)

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

peterdsharpe left a comment

Choose a reason for hiding this comment

Uh oh!

CharlelieLrt commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coreyjadams commented Jul 16, 2026

Uh oh!

CharlelieLrt commented Jul 16, 2026

Uh oh!

coreyjadams commented Jul 16, 2026

Uh oh!

coreyjadams commented Jul 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coreyjadams commented Jul 16, 2026

Uh oh!

coreyjadams commented Jul 17, 2026

Uh oh!

CharlelieLrt left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mnabian left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

coreyjadams commented Jul 16, 2026 •

edited

Loading

greptile-apps Bot commented Jul 16, 2026 •

edited

Loading

CharlelieLrt commented Jul 16, 2026 •

edited

Loading

CharlelieLrt left a comment •

edited

Loading