Update transolver to comply with model standards

[coreyjadams]

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 updates the Transolver model to comply with PhysicsNeMo model implementation standards by adding comprehensive documentation, type annotations, and validation logic.

Major changes:

• Added complete NumPy-style docstrings with proper sections (Parameters, Forward, Outputs, Examples) across all model classes and functions
• Added jaxtyping type annotations for all tensor arguments following MOD-006
• Added input validation with torch.compiler.is_compiling() guards in main forward methods following MOD-005
• Added high-level comments explaining complex tensor operations following MOD-003k
• Updated pyproject.toml to ignore F722 (allows jaxtyping syntax)
• Changed docstring prefixes from """ to r""" for LaTeX compatibility following MOD-003b

Critical issues found:

• MLP and Transolver_block classes inherit from nn.Module instead of physicsnemo.Module, violating MOD-001. These classes should be updated to inherit from physicsnemo.Module to ensure access to serialization, versioning, and registry features.
• Missing input validation in MLP.forward() and Transolver_block.forward() methods (MOD-005 requirement)

Positive aspects:

• Excellent documentation quality with clear examples and cross-references
• Proper use of LaTeX math notation for tensor shapes
• Good high-level comments in complex tensor operations
• Consistent formatting and structure across all files

Important Files Changed

File Analysis

Filename	Score	Overview
physicsnemo/models/transolver/transolver.py	3/5	Added comprehensive docstrings with jaxtyping annotations, input validation, and high-level comments. Found critical issue: MLP and Transolver_block inherit from nn.Module instead of physicsnemo.Module (violates MOD-001).
physicsnemo/models/transolver/Physics_Attention.py	4/5	Added comprehensive docstrings with proper sections, jaxtyping annotations, and input validation. All classes correctly inherit from nn.Module (appropriate for reusable layers per MOD-000a).
physicsnemo/models/transolver/Embedding.py	4/5	Added complete docstrings with proper NumPy-style sections, jaxtyping annotations, LaTeX math notation for tensor shapes, and Examples sections. All classes correctly inherit from nn.Module.
pyproject.toml	5/5	Added F722 to ruff ignore list (allows jaxtyping syntax) and removed trailing whitespace. Changes are appropriate for supporting jaxtyping annotations.

[CharlelieLrt]

Mostly looks good, just a few minor things about formatting and such.
One bigger question though:

• For Physics_Attention.py, could it be moved to physicsnemo/nn, ither with other attention layers in physicsnemo/nn/attention_layers.py, or just in its own module?
• Same question for Embeddings.py

[CharlelieLrt]

LGTM, really nice you could reuse some of the existing layers!

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci