Add point transformer attention

[mnabian]

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.

[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 adds three Point Transformer attention building blocks — LocalPointTransformerBlock, LocalTokenCrossAttentionBlock, and ResidualMLP — implementing k-NN-local grouped vector attention with optional DiT-style AdaLN/AdaLN-Zero conditioning, directly into the production physicsnemo/nn/module/ path.

• All three classes inherit from torch.nn.Module directly, violating MOD-001, which requires all layers in this codebase to use physicsnemo.Module for serialization, versioning, and registry support.
• The classes are placed in production without the required CI test suite (MOD-002a, MOD-008b, MOD-008c): non-regression tests against .pth reference data and .mdlus checkpoint loading tests are absent.
• pos_proj is hardcoded to a 3-dimensional linear input while the documented API accepts arbitrary (N, D) coordinates, causing a silent runtime failure for non-3D inputs.

Important Files Changed

Filename	Overview
physicsnemo/nn/module/point_transformer_attention.py	New file implementing three Point Transformer attention blocks; has multiple standard violations: classes inherit from nn.Module instead of physicsnemo.Module (MOD-001), pos_proj hardcoded to 3D coordinates while API documents arbitrary D, missing forward shape validation (MOD-005), and missing jaxtyping annotations (MOD-006)
test/nn/module/test_point_transformer_attention.py	New test file with good coverage of shapes, gradients, edge cases, and batch masking, but missing required non-regression tests against .pth reference data (MOD-008b) and checkpoint loading tests via physicsnemo.Module.from_checkpoint() (MOD-008c)
physicsnemo/nn/init.py	Adds imports for the three new classes from point_transformer_attention; straightforward and correctly placed
physicsnemo/nn/module/init.py	Re-exports the three new classes; no issues
CHANGELOG.md	Adds a well-written changelog entry describing the new attention blocks; no issues

Comments Outside Diff (1)

1. physicsnemo/nn/module/point_transformer_attention.py, line 568-600 (link)

   self.conditioning is called twice when context_cond is None

   In LocalTokenCrossAttentionBlock.forward, when context_cond is None, kv_source falls back to cond and self.conditioning is called a second time with the exact same input. This is redundant computation that could instead reuse the already-computed result from the first call. Consider caching the full 5-way output and slicing the needed chunks, rather than invoking the MLP a second time.

_{Reviews (1): Last reviewed commit: "add point transformer attention" | Re-trigger Greptile}

[mnabian]

/blossom-ci

[loliverhennigh]

LGTM

[mnabian]

/blossom-ci

[mnabian]

/ok to test 96a60c2

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

/ok to test 9e3582e

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

/ok to test f0fddb8