Revert "Update geotransolver for 2d and 3d use cases"#1619
Merged
coreyjadams merged 1 commit intomainfrom May 5, 2026
Merged
Conversation
This reverts commit b193e03.
Collaborator
Author
|
/blossom-ci |
Contributor
Greptile SummaryThis PR reverts #1502 ("Update geotransolver for 2d and 3d use cases"), removing structured 2D/3D mesh support, the Muon optimizer, multi-model Hydra dispatch, JSONL metrics logging, and the shared free-function refactor (
Important Files Changed
Reviews (1): Last reviewed commit: "Revert "Update geotransolver for 2d and ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+129
to
+184
| if use_te: | ||
| raise ValueError( | ||
| "GALE_FA does not support Transformer Engine backend. " | ||
| "Use use_te=False; TE disables FlashAttention for differing q/k sizes in FLARE attention." | ||
| ) | ||
| super().__init__() | ||
| if state_mixing_mode not in ("weighted", "concat_project"): | ||
| raise ValueError( | ||
| f"Invalid state_mixing_mode: {state_mixing_mode!r}. " | ||
| f"Expected 'weighted' or 'concat_project'." | ||
| ) | ||
| self.state_mixing_mode = state_mixing_mode | ||
| self.use_te = use_te | ||
| self.heads = heads | ||
| self.dim_head = dim_head | ||
| self.scale = 1.0 | ||
| # It is recommended by the FLARE authors to use self.scale = 1 if self.dim_head <= 8 else (self.dim_head ** -0.5) | ||
| # but we use self.scale = 1.0 because the recommended scaling is not tested yet. | ||
| inner_dim = dim_head * heads | ||
|
|
||
| linear_layer = te.Linear if self.use_te else nn.Linear | ||
|
|
||
| # Global queries for FLARE self-attention | ||
| self.q_global = nn.Parameter(torch.randn(1, heads, n_global_queries, dim_head)) | ||
|
|
||
| # Linear projections for self-attention | ||
| self.in_project_x = linear_layer(dim, inner_dim) | ||
| self.self_k = linear_layer(dim_head, dim_head) | ||
| self.self_v = linear_layer(dim_head, dim_head) | ||
|
|
||
| if context_dim > 0: | ||
| # Linear projections for cross-attention | ||
| self.cross_q = linear_layer(dim_head, dim_head) | ||
| self.cross_k = linear_layer(context_dim, dim_head) | ||
| self.cross_v = linear_layer(context_dim, dim_head) | ||
|
|
||
| # Mixing layers for blending self-attention and cross-attention | ||
| if state_mixing_mode == "weighted": | ||
| # Learnable mixing weight between self and cross attention | ||
| self.state_mixing = nn.Parameter(torch.tensor(0.0)) | ||
| else: | ||
| # Concatenate self and cross attention and project back to dim_head | ||
| self.concat_project = nn.Sequential( | ||
| linear_layer(2 * dim_head, dim_head), | ||
| nn.GELU(), | ||
| ) | ||
|
|
||
| # te attention | ||
| if self.use_te: | ||
| self.attn_fn = te.DotProductAttention( | ||
| num_attention_heads=self.heads, | ||
| kv_channels=self.dim_head, | ||
| attention_dropout=dropout, | ||
| qkv_format="bshd", | ||
| softmax_scale=self.scale | ||
| ) |
Contributor
There was a problem hiding this comment.
forward
use_te=True raises ValueError at line 129 before super().__init__() is ever called, so self.use_te is always False. Every subsequent if self.use_te: branch — including self.attn_fn = te.DotProductAttention(...) in __init__ (lines 177–184) and the entire TE blocks inside forward() (lines 237–249, 261–268) — is permanently dead code. This also means the te.Linear option in linear_layer (line 149) can never be selected. Consider removing the dead branches or adding a TODO comment explaining TE support is planned for the future.
| Whether to use learned concrete dropout instead of standard dropout. | ||
| Default is ``False``. | ||
| state_mixing_mode : str, optional | ||
| How to blend self-attention and cross-attention outputs. ``"weighted"`` uses |
Contributor
There was a problem hiding this comment.
gale.py for state_mixing_mode — extraneous spaces between the period and the first backtick-quoted option.
Suggested change
| How to blend self-attention and cross-attention outputs. ``"weighted"`` uses | |
| How to blend self-attention and cross-attention outputs. ``"weighted"`` uses |
Collaborator
Author
|
None of the greptile comments will be resolved because we are reverting here. |
Collaborator
|
This PR fixes the issue with the CFD checkpoints |
ktangsali
approved these changes
May 5, 2026
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
2 participants
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.
Reverts #1502