Align NAT functionals with FunctionSpec by loliverhennigh · Pull Request #1615 · NVIDIA/physicsnemo

loliverhennigh · 2026-05-04T23:17:07Z

PhysicsNeMo Pull Request

Description

Cleans up some of the NAT functional stuff and others that made it past my radar.

copy-pr-bot · 2026-05-04T23:17:11Z

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 · 2026-05-04T23:20:19Z

Greptile Summary

This PR wraps the existing na1d/na2d/na3d NATTEN functions in FunctionSpec-derived classes (NeighborhoodAttention1D/2D/3D), adds them to the benchmark registry alongside two geometry specs, ships a new doc page, and updates FNC-000 to officially allow lightweight pure-PyTorch helpers to stay as plain functions.

P1 – test will fail when natten is installed: test_natten_version_mismatch_error_message monkeypatches get_installed_version but available is frozen at class-definition time; with natten >= 0.21.5 present, dispatch succeeds and pytest.raises(ImportError) raises Failed: DID NOT RAISE.
P2 – unreachable raise in _NeighborhoodAttention.dispatch: the final bare re-raise is only reachable if get_installed_version returns None yet _natten.functional does not raise, a combination that cannot occur with OptionalImport; worth a clarifying comment or a more descriptive fallback error.

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/natten.py	Wraps na1d/na2d/na3d in FunctionSpec subclasses with a custom dispatch that enriches ImportError messages; the dispatch error-handling has an unreachable code path and an associated test that breaks in environments with a conforming natten install.
test/nn/functional/test_natten.py	Adds FunctionSpec contract tests and error-message tests; test_natten_version_mismatch_error_message is unreliable when natten >= 0.21.5 is installed because available=True is fixed at class-definition time.
benchmarks/physicsnemo/nn/functional/registry.py	Adds NeighborhoodAttention1D/2D/3D and geometry specs to FUNCTIONAL_SPECS; straightforward and correct.
CODING_STANDARDS/FUNCTIONAL_APIS.md	Clarifies FNC-000 to allow lightweight pure-PyTorch helpers to stay as plain functions; well-reasoned criteria and rationale.
docs/api/nn/functionals/neighborhood_attention.rst	New RST page auto-documenting na1d/na2d/na3d from the correct public path; no issues.
docs/api/physicsnemo.nn.functionals.rst	Adds the new neighborhood_attention page to the functionals TOC; no issues.

_{Reviews (1): Last reviewed commit: "Align NAT functionals with FunctionSpec" | Re-trigger Greptile}

greptile-apps · 2026-05-04T23:20:24Z

+    @classmethod
+    def dispatch(cls, *args, **kwargs):
+        try:
+            return super().dispatch(*args, **kwargs)
+        except ImportError as exc:
+            installed_version = get_installed_version("natten")
+            if installed_version is not None:
+                raise ImportError(
+                    f"{cls._NATTEN_REQUIREMENT} is required for {cls.__name__}, "
+                    f"but found natten {installed_version}"
+                ) from exc
+            try:
+                _natten.functional
+            except ImportError as missing_exc:
+                raise missing_exc from exc
+            raise


Unreachable raise at the end of the dispatch override

The final bare raise is only reached when installed_version is None and _natten.functional does not raise ImportError. Because installed_version is None means the package metadata is absent and the OptionalImport proxy always raises ImportError on attribute access when the underlying module cannot be imported, this branch is currently unreachable in practice. If _natten.functional somehow succeeds (e.g., metadata is absent but the package is importable), re-raising exc produces the cryptic "No available implementations found" message rather than any natten-specific hint. The code path should either be documented with a comment explaining the invariant, or the final raise should emit a more descriptive error.

pzharrington · 2026-05-04T23:52:09Z

@loliverhennigh my main question is have you tested that this still works with ShardTensor and domain parallelism? Want to make sure the dispatch path in that case does the right thing. Otherwise looks good

loliverhennigh · 2026-05-05T17:28:35Z

-from physicsnemo.nn.functional.geometry import SignedDistanceField
+from physicsnemo.nn.functional.geometry import (
+    MeshPoissonDiskSample,
+    MeshToVoxelFraction,


Some of the functionals got left out so adding them back in

loliverhennigh · 2026-05-05T17:32:54Z

-All functionals must be implemented with `FunctionSpec`, even if only a single
-implementation exists. This ensures the operation participates in validation
-and benchmarking through input generators and `compare_forward` (and
+All functionals with backend dispatch, optional accelerated implementations, or


Basically going to make some space for if we have very simple python functions and dont want to implement the whole FunctionSpec stuff

megnvidia

tiny suggestion

megnvidia · 2026-05-05T21:00:43Z

rst files looked good to me, I had one very minor suggestion for Nvidia style guide rules. Looks like you still have other comments, but I am ready to approve of this when you need.

loliverhennigh · 2026-05-05T23:00:33Z

/blossom-ci

loliverhennigh · 2026-05-06T18:48:14Z

/multi-gpu-ci

loliverhennigh · 2026-05-06T18:52:38Z

/blossom-ci

coreyjadams · 2026-05-06T19:24:30Z

/multi-gpu-ci

I have no idea if this still does anything. You might need to test manually on a cluster.

loliverhennigh · 2026-05-07T21:30:02Z

Hey @coreyjadams and @pzharrington , Just manually ran on two GB200s and it passed. I think we are all good to merge unless there is something else

peterdsharpe · 2026-05-07T21:32:05Z

-            DualInteractionPlan,
-            SourceAggregates,
-        )
+        from physicsnemo.experimental.models.globe.cluster_tree import ClusterTree


Motivation for this edit? Seems unrelated

loliverhennigh · 2026-05-07T21:33:28Z

/blossom-ci

peterdsharpe · 2026-05-07T21:35:13Z

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .ops import (


This change of equivariant_ops -> equivariant.ops is backwards-incompatible. We should at-bare-minimum have a CHANGELOG entry for this, but realistically this probably calls for a larger discussion about whether this is worth breaking back-compat for, given that there may be downstream users that are consumers of this API surface. In general, we shouldn't be breaking public (i.e., non-_-prefixed) API surface unless it's exceedingly important that we do so.

Also, tagging @laserkelvin for API thoughts on the equivariant module (not sure if you have any uses of this?)

I don't think I have an issue with naming alignment itself (i.e. nn.functional.equivariant) - I'm just confused why this is happening now, and to your (@peterdsharpe's) point this is backwards incompatible. AFAIK there aren't public consumers of that API yet, but please make sure that the tests associated with this change pass

Unless we introduced equivariant_ops within this current release cycle (in which case it's not part of the public API yet)? What does the git blame say?

It's not in the 2.0.0 release, and I think you added it 2 months ago

Ah - if it wasn't ever released publicly in an official release cut, then I'm more ok with breaking this API.

That said, it's not obvious to me that either naming convention is better than the other. In the absence of motivation, defaulting to the existing API seems like the best path forward. That said, there may be motivation for this change that I'm not seeing - @loliverhennigh can you speak to the motivation for this change?

External customers should only use functionals via physicsnemo.nn.functionals.xyz. This is how the coding standards and docs are laid out at least. This PR does not change that import so I think we are good. Only internal imports might change.

loliverhennigh · 2026-05-07T23:22:29Z

/blossom-ci

megnvidia

lgtm

peterdsharpe · 2026-05-08T18:13:02Z

@@ -22,11 +22,15 @@


 @overload
-def smooth_log(x: Float[torch.Tensor, "..."]) -> Float[torch.Tensor, "..."]: ...
+def smooth_log(x: Float[torch.Tensor, "..."]) -> Float[torch.Tensor, "..."]:
+    """Apply smooth log elementwise to a tensor."""


These are @overloads - are docstrings needed to satisfy interrogate? I would hope that they would not be needed, but I don't know if interrogate is smart enough to handle this.

Maybe interrogate has an option to ignore overloaded functions?

peterdsharpe · 2026-05-08T18:21:35Z

@loliverhennigh can we flesh out the discussion in unresolved comment threads before merging?

loliverhennigh · 2026-05-08T18:41:42Z

/blossom-ci

Co-authored-by: megnvidia <mmiranda@nvidia.com>

loliverhennigh · 2026-05-08T22:15:33Z

/blossom-ci

loliverhennigh requested review from ktangsali, megnvidia and pzharrington as code owners May 4, 2026 23:17

greptile-apps Bot reviewed May 4, 2026

View reviewed changes

pzharrington reviewed May 4, 2026

View reviewed changes

Comment thread physicsnemo/nn/functional/attention/neighborhood_attention.py

loliverhennigh commented May 5, 2026

View reviewed changes

loliverhennigh force-pushed the functional-api-cleanup branch from 8da5e27 to 280d401 Compare May 5, 2026 18:53

loliverhennigh requested review from coreyjadams and peterdsharpe as code owners May 5, 2026 18:53

loliverhennigh force-pushed the functional-api-cleanup branch 2 times, most recently from f0cb8a9 to e3847e8 Compare May 5, 2026 19:08

loliverhennigh requested a review from CharlelieLrt as a code owner May 5, 2026 19:08

megnvidia reviewed May 5, 2026

View reviewed changes

Comment thread CODING_STANDARDS/FUNCTIONAL_APIS.md Outdated

loliverhennigh force-pushed the functional-api-cleanup branch 2 times, most recently from 297f9cb to 62a0ef7 Compare May 5, 2026 22:41

peterdsharpe reviewed May 7, 2026

View reviewed changes

pzharrington approved these changes May 7, 2026

View reviewed changes

megnvidia approved these changes May 8, 2026

View reviewed changes

peterdsharpe reviewed May 8, 2026

View reviewed changes

loliverhennigh force-pushed the functional-api-cleanup branch from 8519251 to 6885e86 Compare May 8, 2026 18:41

peterdsharpe self-requested a review May 8, 2026 19:33

loliverhennigh and others added 7 commits May 8, 2026 15:03

Align NAT functionals with FunctionSpec

2dff29e

Apply suggestions from code review

cdd413f

Co-authored-by: megnvidia <mmiranda@nvidia.com>

Add NAT FunctionSpec docstrings

869d3fa

Add equivariant functional overload docstrings

0f48fde

Stabilize NAT version mismatch test

f1476d3

Address functional API review comments

5e30322

Restore field kernel local imports

e638afc

loliverhennigh force-pushed the functional-api-cleanup branch from b8d31d3 to e638afc Compare May 8, 2026 22:04

loliverhennigh added 3 commits May 8, 2026 15:07

Keep equivariant ops module path

0b1a684

Ignore overload stubs in docstring coverage

8a7659c

Keep docstring coverage policy unchanged

5d729e9

Conversation

loliverhennigh commented May 4, 2026

PhysicsNeMo Pull Request

Description

Uh oh!

copy-pr-bot Bot commented May 4, 2026

Uh oh!

greptile-apps Bot commented May 4, 2026

Greptile Summary

Important Files Changed

Uh oh!

Uh oh!

greptile-apps Bot May 4, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pzharrington commented May 4, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

megnvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

megnvidia commented May 5, 2026

Uh oh!

loliverhennigh commented May 5, 2026

Uh oh!

loliverhennigh commented May 6, 2026

Uh oh!

loliverhennigh commented May 6, 2026

Uh oh!

coreyjadams commented May 6, 2026

Uh oh!

loliverhennigh commented May 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

loliverhennigh commented May 7, 2026

Uh oh!

peterdsharpe May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peterdsharpe May 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

loliverhennigh commented May 7, 2026

Uh oh!

megnvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peterdsharpe commented May 8, 2026

Uh oh!

loliverhennigh commented May 8, 2026

Uh oh!

loliverhennigh commented May 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

peterdsharpe May 7, 2026 •

edited

Loading

peterdsharpe May 8, 2026 •

edited

Loading