GeoFLARE: added GALE_FA, an alternate attention to GALE, for GeoTransolver

[dakhare-creator]

PhysicsNeMo Pull Request

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.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds GAFLARE (Geometry-Aware FLARE), an alternative attention mechanism to GALE for the GeoTransolver model.

Key Changes:

• New gaflare.py module implementing FLARE-based attention with geometry-aware context
• Modified gale.py to support dynamic attention type selection via attention_type parameter
• Updated geotransolver.py to expose attention_type parameter (default: "GALE")

Critical Issues Found:

• gaflare.py line 121: use_te parameter is hardcoded to False, ignoring the constructor argument
• gaflare.py line 122: Attention scale hardcoded to 1.0 instead of proper dim_head**-0.5 scaling
• gale.py line 393: Unsafe use of globals() for class selection - security risk and violates MOD-009 coding standard

Minor Issues:

• Typo in docstring: "Origional" → "Original"
• Missing newline at end of gaflare.py

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/geotransolver/gaflare.py	New GAFLARE attention mechanism added. Critical issues: use_te parameter hardcoded to False (line 121), attention scale set to 1.0 instead of proper scaling. Minor: typo in docstring, missing newline at EOF.
physicsnemo/experimental/models/geotransolver/gale.py	Integration changes to support attention type selection. Critical security issue: unsafe use of globals() for dynamic class instantiation (line 393) allows potential code execution.
physicsnemo/experimental/models/geotransolver/geotransolver.py	Added attention_type parameter to allow switching between GALE and GAFLARE. Clean pass-through implementation with proper default value.

[greptile-apps]

_{3 files reviewed, 6 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

physicsnemo/experimental/models/geotransolver/gale.py
Using globals() for dynamic class selection is unsafe - allows arbitrary code execution if attention_type comes from untrusted input. This violates MOD-009 (avoid string-based class selection). Consider explicit mapping:

ATTENTION_CLASSES = {
    "GALE": GALE,
    "GAFLARE": GAFLARE,
}
if attention_type not in ATTENTION_CLASSES:
    raise ValueError(f"Unknown attention_type: {attention_type}")
self.Attn = ATTENTION_CLASSES[attention_type](
    hidden_dim,
    heads=num_heads,
    # ... rest of arguments
)

[coreyjadams]

Hi @dakhare-creator - thanks for opening this. Overall it seems reasonable. Good to put it in experimental. I made a couple comments in the code, but also:

• it'd be unusual for use to publish a "GeoFlare" model without also making available "Flare". Can you add standard FLARE attention to the experimental/nn folder so we can do that too?
• We can also probably easily create a flare.py itself for the model.
• How much overlap in the math and attention blocks is there between PhysicsAttention and FLARE Attention? PhysicsAttention can be used on 2d and 3d data, could this? I strongly suspect yes. Can we make this modular so we can reuse as much as possible between attention layers?

Once we're ready, let's also invite the FLARE authors to take a look.

[vpuri3]

Hi @dakhare-creator, I've left several comments in the code. Implementing FLARE straight on a geom problem can be error prone. I'd recommend first validating FLARE on one of the paper's benchmark problems

https://github.com/vpuri3/FLARE.py/blob/master/pdebench/models/flare.py

and then exposing the relevant layers to the geometry solvers interface.

[mnabian]

/blossom-ci

[coreyjadams]

Hi @dakhare-creator - with the geometry encoded FLARE, does this also enable standard FLARE attention and model?

[dakhare-creator]

Hi @dakhare-creator - with the geometry encoded FLARE, does this also enable standard FLARE attention and model?

Hi @coreyjadams, Yes if the context is None, the attention mechanism will be standard FLARE.

[mnabian]

/blossom-ci

[vpuri3]

Hi all, is there any update on this?

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci