Fix Transolver structured-embedding flattening

[peterdsharpe]

PhysicsNeMo Pull Request

Transolver.forward crashed when given a spatially-shaped embedding for structured data with unified_pos=False. This corrects the reshape so the embedding is flattened to tokens the same way fx already is.

Why

For structured data the forward pass flattens fx from (B, H, W[, D], C) to (B, N, C), then concatenates it with embedding. The embedding branch did the opposite: it reshaped a non-3D embedding to (B, *structured_shape, -1), leaving it 4D/5D. The next line, torch.cat((embedding, fx), -1), requires equal rank, so it always raised RuntimeError: Tensors must have same number of dimensions.

The branch could never succeed, and it contradicted the docstring, which told users the embedding "should match fx spatial dimensions" (exactly the shape that crashed). The fix flattens the embedding consistently with fx:

# before: un-flattens to (B, H, W[, D], C_emb), then cat fails
-embedding = embedding.reshape(embedding.shape[0], *self.structured_shape, -1)
# after: flattens to (B, N, C_emb), aligned per-token with fx
+embedding = embedding.reshape(embedding.shape[0], -1, embedding.shape[-1])

What it affects

• Transolver and FLARE (which inherits Transolver.forward) now accept spatially-shaped embeddings for structured data. GeoTransolver has its own forward and is unaffected.
• No change for existing callers. Every current caller passes a flattened (B, N, C_emb) embedding, uses unified_pos=True (auto 3D buffer), or runs unstructured (structured_shape=None), so the branch was skipped. The change only turns the previously-crashing path into correct output, so it cannot regress working code. No signature or return-type changes.

Tests

Added a parametrized (2D and 3D) regression test that runs the structured, unified_pos=False path with a spatially-shaped embedding, checks the output shape, and confirms it matches the equivalent pre-flattened (B, N, C_emb) input. Verified that irregular, unified_pos, and pre-flattened structured paths are unchanged.

Docs

Clarified the embedding parameter docstring: for structured data it may be passed either flattened as (B, N, C_emb) or with the same spatial layout as fx; the output follows fx's layout.

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.

[copy-pr-bot]

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]

Greptile Summary

This PR fixes a one-line bug in Transolver.forward where a spatially-shaped embedding was reshaped back to (B, *structured_shape, C_emb) instead of being flattened to (B, N, C_emb), causing every call through that branch to raise RuntimeError from a rank mismatch in torch.cat. The fix aligns the embedding reshape with the already-correct fx flatten, and adds docstring clarifications.

• Bug fix: embedding.reshape(embedding.shape[0], *self.structured_shape, -1) → embedding.reshape(embedding.shape[0], -1, embedding.shape[-1]), making the branch reachable and correct for both 2D and 3D structured inputs.
• Regression test: Parametrized over 2D and 3D spatial shapes, verifies output shape and that spatially-shaped and pre-flattened embeddings produce bit-identical results (atol=1e-6).

Important Files Changed

Filename	Overview
physicsnemo/models/transolver/transolver.py	Fixes incorrect reshape in the structured-embedding flattening branch: changes the reshape from (B, *structured_shape, -1) (which left embedding 4D/5D, mismatching flattened fx) to (B, -1, embedding.shape[-1]) (correctly producing (B, N, C_emb)). Docstrings updated accordingly.
test/models/transolver/test_transolver.py	Adds a parametrized regression test covering both 2D and 3D structured cases with spatially-shaped embeddings; verifies output shape and that spatial and pre-flattened inputs yield identical results.

_{Reviews (1): Last reviewed commit: "Enhance Transolver documentation and add..." | Re-trigger Greptile}

[peterdsharpe]

/blossom-ci

[peterdsharpe]

/ok to test 95892e0

[coreyjadams]

Can we merge this into the 2D-3D cleanup PR?

[coreyjadams]

Specifically, I think we could merge this into here perhaps: #1676

[peterdsharpe]

Weak preference to keep these separate (since #1676 seems to have bigger scope, and in-general I'm a fan of more-and-smaller PRs rather than fewer-and-larger PRs). I'm comfortable waiting until #1676 merges before merging this in, if you prefer. But, if you feel strongly that we should combine these PRs, that's fine too! Will sync up offline with you!