From 06e7e5519972765ac98704f07834c1a5e77d2664 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Fri, 16 Oct 2026 10:27:16 -0700
Subject: [PATCH 1/8] na1d/na2d/na3d: dispatch ShardTensor handler under
 torch.compile

Dynamo folds has_torch_function to False for traceable wrapper subclasses,
so a compiled na2d called the raw kernel on the ShardTensor and skipped the
halo exchange (wrong results, or rank-divergent collectives). Under compile,
route subclass calls through an allow_in_graph helper: Dynamo records one
graph node and fake propagation runs the handler, as for torch ops.

Test uses a pure-torch reference kernel, so it runs without NATTEN.
---
 physicsnemo/nn/functional/natten.py           |  42 +++++++
 .../ops/test_natten_compile.py                | 112 ++++++++++++++++++
 2 files changed, 154 insertions(+)
 create mode 100644 test/domain_parallel/ops/test_natten_compile.py

diff --git a/physicsnemo/nn/functional/natten.py b/physicsnemo/nn/functional/natten.py
index 225acf33d3..69ad827730 100644
--- a/physicsnemo/nn/functional/natten.py
+++ b/physicsnemo/nn/functional/natten.py
@@ -26,6 +26,40 @@
 _natten = OptionalImport("natten")
 
 
+def _subclass_under_compile(*tensors: torch.Tensor) -> bool:
+    r"""True while Dynamo traces a call with a tensor-subclass operand.
+
+    Dynamo folds ``has_torch_function`` to ``False`` for traceable wrapper
+    subclasses such as ``ShardTensor``, so the check above alone would hand the
+    subclass to the raw kernel. ``type()`` still reports the subclass.
+    """
+    return torch.compiler.is_compiling() and any(
+        type(t) is not torch.Tensor for t in tensors
+    )
+
+
+@torch._dynamo.allow_in_graph
+def _dispatch_in_graph(
+    name: str,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kernel_size: int,
+    dilation: int,
+    kwargs: dict[str, Any],
+) -> torch.Tensor:
+    r"""Subclass dispatch as one opaque graph node.
+
+    Dynamo records the call and runs it under fake propagation, where
+    ``__torch_function__`` dispatches to the subclass handler (for
+    ``ShardTensor``, the halo-exchange path), as for ``torch`` ops.
+    """
+    func = _NA_FUNCS[name]
+    return handle_torch_function(
+        func, (q, k, v), q, k, v, kernel_size, dilation=dilation, **kwargs
+    )
+
+
 def na1d(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -73,6 +107,8 @@ def na1d(
             dilation=dilation,
             **kwargs,
         )
+    if _subclass_under_compile(q, k, v):
+        return _dispatch_in_graph("na1d", q, k, v, kernel_size, dilation, kwargs)
     return _natten.functional.na1d(q, k, v, kernel_size, dilation=dilation, **kwargs)
 
 
@@ -123,6 +159,8 @@ def na2d(
             dilation=dilation,
             **kwargs,
         )
+    if _subclass_under_compile(q, k, v):
+        return _dispatch_in_graph("na2d", q, k, v, kernel_size, dilation, kwargs)
     return _natten.functional.na2d(q, k, v, kernel_size, dilation=dilation, **kwargs)
 
 
@@ -173,7 +211,11 @@ def na3d(
             dilation=dilation,
             **kwargs,
         )
+    if _subclass_under_compile(q, k, v):
+        return _dispatch_in_graph("na3d", q, k, v, kernel_size, dilation, kwargs)
     return _natten.functional.na3d(q, k, v, kernel_size, dilation=dilation, **kwargs)
 
 
+_NA_FUNCS = {"na1d": na1d, "na2d": na2d, "na3d": na3d}
+
 __all__ = ["na1d", "na2d", "na3d"]
diff --git a/test/domain_parallel/ops/test_natten_compile.py b/test/domain_parallel/ops/test_natten_compile.py
new file mode 100644
index 0000000000..a3d046d1d2
--- /dev/null
+++ b/test/domain_parallel/ops/test_natten_compile.py
@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""``torch.compile`` over sharded neighborhood attention (issue #1986).
+
+Dynamo treats ``ShardTensor`` as a traceable wrapper subclass and folds
+``has_torch_function`` to ``False`` for it, so a traced ``na2d`` would call the
+raw kernel on the ShardTensor instead of the halo-exchanging handler. The raw
+kernel here is a pure-torch reference that, like an opaque kernel, refuses
+distributed inputs; it also lets the test run without NATTEN installed.
+"""
+
+import pytest
+import torch
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import Shard
+
+from physicsnemo.distributed import DistributedManager
+from physicsnemo.domain_parallel import ShardTensor, scatter_tensor
+
+from .utils import collective_assert_close, sharded_to_local
+
+
+def _na2d_reference(q, k, v, kernel_size, dilation=1, **_):
+    r"""NATTEN-semantics 2D neighborhood attention on plain ``(B, H, W, h, D)``
+    tensors: each query attends to the ``kernel_size`` window clamped to the
+    image, not zero-padded."""
+    if isinstance(q, (ShardTensor, DTensor)):
+        raise RuntimeError("raw neighborhood-attention kernel got a distributed tensor")
+    assert dilation == 1
+    r = kernel_size
+
+    def window(n):
+        start = torch.arange(n, device=q.device).clamp(r // 2, n - 1 - r // 2) - r // 2
+        return start[:, None] + torch.arange(r, device=q.device)
+
+    ih, iw = window(q.shape[1]), window(q.shape[2])
+    kk = k[:, ih][:, :, :, iw]  # (B, H, r, W, r, h, D)
+    vv = v[:, ih][:, :, :, iw]
+    logits = torch.einsum("bxyhd,bxiyjhd->bxyhij", q, kk) * q.shape[-1] ** -0.5
+    p = logits.flatten(-2).softmax(-1).view_as(logits)
+    return torch.einsum("bxyhij,bxiyjhd->bxyhd", p, vv)
+
+
+@pytest.fixture
+def natten_reference(monkeypatch):
+    r"""Route the raw NATTEN call (eager and sharded paths) to the reference."""
+    from types import SimpleNamespace
+
+    from physicsnemo.domain_parallel.shard_utils import natten_patches
+    from physicsnemo.nn.functional import natten as natten_functional
+
+    stub = SimpleNamespace(functional=SimpleNamespace(na2d=_na2d_reference))
+    monkeypatch.setattr(natten_functional, "_natten", stub)
+    monkeypatch.setattr(natten_patches, "_natten", stub)
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.timeout(300)
+@pytest.mark.parametrize("backward", [False, True])
+def test_na2d_compiled_on_shard_tensor(distributed_mesh, natten_reference, backward):
+    r"""Compiled ``na2d`` on height-sharded q/k/v runs the halo-exchange handler
+    and matches the single-device result, forward and backward."""
+    from physicsnemo.nn.functional.natten import na2d
+
+    dm = DistributedManager()
+    kernel_size = 5
+    shape = (1, 8 * distributed_mesh.size(0), 12, 2, 8)  # (B, H, W, heads, D)
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(shape, device=dm.device) for _ in range(3))
+    sq, sk, sv = (
+        scatter_tensor(t, 0, distributed_mesh, (Shard(1),), requires_grad=backward)
+        for t in (q, k, v)
+    )
+
+    torch._dynamo.reset()
+    compiled = torch.compile(
+        lambda q, k, v: na2d(q, k, v, kernel_size=kernel_size),
+        backend="aot_eager",
+        dynamic=False,
+        fullgraph=True,  # the handler must be a graph node, not a graph break
+    )
+    out = compiled(sq, sk, sv)
+
+    if backward:
+        q, k, v = (t.detach().requires_grad_(True) for t in (q, k, v))
+    ref = _na2d_reference(q, k, v, kernel_size)
+
+    assert isinstance(out, ShardTensor)
+    assert out._spec.placements == (Shard(1),)
+    collective_assert_close(ref, sharded_to_local(out), atol=1e-4, rtol=1e-4, msg="out")
+
+    if backward:
+        out.square().sum().backward()
+        ref.square().sum().backward()
+        for name, st, t in (("q", sq, q), ("k", sk, k), ("v", sv, v)):
+            collective_assert_close(
+                t.grad, sharded_to_local(st.grad), atol=1e-3, rtol=1e-3, msg=name
+            )

From 3683ce797f63cd62232f76a94c3d2dbc68dbc1a2 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Fri, 16 Oct 2026 10:27:17 -0700
Subject: [PATCH 2/8] ShardTensor select/select_backward: normalize negative
 dim

The dispatch handlers compared shard dims against dim as written: dim=-1
(stack's backward, e.g. RoPE rotate-half under compile) relabeled Shard(2)
as Shard(1) on the gradient, and a negative dim on the sharded axis escaped
the guard.
---
 .../domain_parallel/shard_utils/index_ops.py  |  3 +
 test/domain_parallel/ops/test_select.py       | 64 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/physicsnemo/domain_parallel/shard_utils/index_ops.py b/physicsnemo/domain_parallel/shard_utils/index_ops.py
index b7d5d305ff..f275a5e2f7 100644
--- a/physicsnemo/domain_parallel/shard_utils/index_ops.py
+++ b/physicsnemo/domain_parallel/shard_utils/index_ops.py
@@ -884,6 +884,8 @@ def sharded_select_helper(tensor: ShardTensor, dim: int, index: int) -> ShardTen
     # If it's along an unsharded dimension, there is nearly nothing to do.
 
     input_spec = tensor._spec
+    # Dispatch passes ``dim`` as written (e.g. -1 from stack's backward).
+    dim = dim % len(input_spec.shape)
 
     input_placements = input_spec.placements
 
@@ -985,6 +987,7 @@ def sharded_select_backward_helper(
     # If it's along an unsharded dimension, there is nearly nothing to do.
 
     input_placements = grad_output._spec.placements
+    dim = dim % len(input_sizes)
 
     output_stride = _stride_from_contiguous_shape_C_style(input_sizes)
 
diff --git a/test/domain_parallel/ops/test_select.py b/test/domain_parallel/ops/test_select.py
index f675221f96..87be2dae3f 100644
--- a/test/domain_parallel/ops/test_select.py
+++ b/test/domain_parallel/ops/test_select.py
@@ -302,3 +302,67 @@ def test_index_select_method_spelling(distributed_mesh):
         {},
         check_grads=True,
     )
+
+
+# ---------------------------------------------------------------------------
+# aten.select / select_backward with a negative dim (dispatch-level handlers)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.parametrize("dim", [-1, 4])
+def test_select_negative_dim_keeps_shard_dim(distributed_mesh, dim):
+    r"""The dispatch handlers see ``dim`` as given: a negative dim must not
+    shift the shard dim (``Shard(2)`` stays ``Shard(2)`` when the last dim goes),
+    and select_backward must insert the size at the same position."""
+    dm = DistributedManager()
+    full = torch.randn(2, 3, 8 * distributed_mesh.size(0), 5, 4, device=dm.device)
+    st = scatter_tensor(full, 0, distributed_mesh, (Shard(2),))
+
+    with torch._C.DisableTorchFunctionSubclass():
+        out = torch.ops.aten.select.int(st, dim, 1)
+        back = torch.ops.aten.select_backward.default(out, list(full.shape), dim, 1)
+
+    assert out._spec.placements == (Shard(2),)
+    assert tuple(out.shape) == (2, 3, full.shape[2], 5)
+    torch.testing.assert_close(out.full_tensor(), full.select(dim, 1))
+    assert back._spec.placements == (Shard(2),)
+    torch.testing.assert_close(
+        back.full_tensor(),
+        torch.ops.aten.select_backward.default(
+            full.select(dim, 1), list(full.shape), dim, 1
+        ),
+    )
+
+
+@pytest.mark.multigpu_static
+def test_select_negative_dim_on_shard_dim_rejected(distributed_mesh):
+    r"""Selecting the sharded dim through a negative index is rejected too."""
+    from physicsnemo.domain_parallel.shard_utils.patch_core import MissingShardPatch
+
+    dm = DistributedManager()
+    full = torch.randn(2, 8 * distributed_mesh.size(0), device=dm.device)
+    st = scatter_tensor(full, 0, distributed_mesh, (Shard(1),))
+    with torch._C.DisableTorchFunctionSubclass(), pytest.raises(MissingShardPatch):
+        torch.ops.aten.select.int(st, -1, 0)
+
+
+@pytest.mark.multigpu_static
+def test_compiled_stack_backward_keeps_shard_dim(distributed_mesh):
+    r"""``torch.stack(..., dim=-1)`` (RoPE's rotate-half) under compile: its
+    backward selects with ``dim=-1``; the input grad must stay ``Shard(2)``."""
+    dm = DistributedManager()
+    full = torch.randn(1, 4, 8 * distributed_mesh.size(0), 6, 8, device=dm.device)
+    st = scatter_tensor(full, 0, distributed_mesh, (Shard(2),), requires_grad=True)
+
+    def rotate_half(x):
+        return torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2)
+
+    torch._dynamo.reset()
+    out = torch.compile(rotate_half, backend="aot_eager", dynamic=False)(st)
+    out.square().sum().backward()
+
+    ref = full.clone().requires_grad_(True)
+    rotate_half(ref).square().sum().backward()
+    assert st.grad._spec.placements == (Shard(2),)
+    torch.testing.assert_close(st.grad.full_tensor(), ref.grad)

From ab47a2b539e019f3ac1b3780a05fcda9b434095d Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:42:24 -0700
Subject: [PATCH 3/8] ShardTensorSpec: tolerate symbolic shapes under
 torch.compile

With dynamic shapes, tensor_meta holds SymInts, which are unhashable: the
spec hash (DTensor sharding propagation) and validate_aligned_sharding's
layout key both raised. The hash falls back to (mesh, placements); the
alignment check skips symbolic layouts.
---
 .../domain_parallel/_shard_tensor_spec.py     | 10 +++++--
 .../test_spec_and_plan_cache.py               | 27 +++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/physicsnemo/domain_parallel/_shard_tensor_spec.py b/physicsnemo/domain_parallel/_shard_tensor_spec.py
index 72815ce9ca..e5b5e23a52 100644
--- a/physicsnemo/domain_parallel/_shard_tensor_spec.py
+++ b/physicsnemo/domain_parallel/_shard_tensor_spec.py
@@ -98,8 +98,12 @@ def _hash_impl(self) -> int:
             hash_items.append(self.tensor_meta.dtype)
         if self._sharding_shapes is not None:
             hash_items.append(tuple(sorted(self._sharding_shapes.items())))
-        hash_tuple = tuple(hash_items)
-        return hash(hash_tuple)
+        try:
+            return hash(tuple(hash_items))
+        except TypeError:
+            # Symbolic (SymInt) shapes under torch.compile are unhashable;
+            # collisions are allowed and __eq__ still discriminates.
+            return hash((self.mesh, self.placements))
 
     def __setattr__(self, attr: str, value) -> None:
         r"""Invalidate the cached hash when the sharding shapes change.
@@ -731,6 +735,8 @@ def validate_aligned_sharding(specs, op_name: str) -> None:
             continue
         cached = getattr(spec, "_sharding_shapes", None) or {}
         shape = tuple(spec.tensor_meta.shape)
+        if any(type(d) is not int for d in shape):
+            continue  # symbolic shapes: boundaries can't be compared here
         key = (spec.mesh, len(shape), placements, tuple(shape[d] for _, d in sharded))
         known = reference.setdefault(key, {})
         for mesh_dim, tensor_dim in sharded:
diff --git a/test/domain_parallel/test_spec_and_plan_cache.py b/test/domain_parallel/test_spec_and_plan_cache.py
index 7e7346c7e7..4cc2cd6c01 100644
--- a/test/domain_parallel/test_spec_and_plan_cache.py
+++ b/test/domain_parallel/test_spec_and_plan_cache.py
@@ -238,3 +238,30 @@ def test_record_consumer_stream_unwraps_without_autograd(distributed_mesh, monke
     (tensor,) = recorded
     assert tensor is st._local_tensor
     assert tensor.grad_fn is None
+
+
+# ---------------------------------------------------------------------------
+# Symbolic shapes (torch.compile dynamic=True)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.timeout(120)
+def test_compile_dynamic_non_leaf(distributed_mesh):
+    r"""A non-leaf ShardTensor compiled with symbolic shapes, fwd + bwd.
+
+    Fakeifying it routes through DTensor sharding propagation, which hashes
+    the spec, and the fallback's alignment check keys a dict on the shape:
+    both see ``SymInt`` dims and must not require them to be hashable.
+    """
+    st = shard_tensor_factory(distributed_mesh, requires_grad=True, uneven=True)
+    y = st * 1.0  # non-leaf, created eagerly
+
+    torch._dynamo.reset()
+    out = torch.compile(lambda t: t.sin(), backend="aot_eager", dynamic=True)(y)
+    out.sum().backward()
+
+    torch.testing.assert_close(out.to_local(), y.to_local().sin())
+    torch.testing.assert_close(
+        st.grad.to_local(), st.to_local().detach().cos(), check_device=False
+    )

From f224b9178bab0ea12037e101e6b48b45f24c967e Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Wed, 21 Oct 2026 10:03:28 -0700
Subject: [PATCH 4/8] ShardTensor: route mixed DTensor operands through the
 fallback under compile

Under tracing, unpatched ops went straight to dispatch; a plain DTensor
operand (e.g. a Replicate parameter like the DiT mask token) then took the
op and rejected the ShardTensor ('got mixed torch.Tensor and DTensor').
Such ops now use the DTensor fallback, as in eager.
---
 physicsnemo/domain_parallel/shard_tensor.py | 12 +++++++-
 test/domain_parallel/test_compile.py        | 33 +++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/physicsnemo/domain_parallel/shard_tensor.py b/physicsnemo/domain_parallel/shard_tensor.py
index ea4209968c..1c528fbdab 100644
--- a/physicsnemo/domain_parallel/shard_tensor.py
+++ b/physicsnemo/domain_parallel/shard_tensor.py
@@ -797,6 +797,14 @@ def _is_tracing(args: object, kwargs: object = None) -> bool:
     return False
 
 
+def _has_plain_dtensor(args: object, kwargs: object = None) -> bool:
+    r"""True when any operand is a ``DTensor`` that is not a ``ShardTensor``."""
+    return any(
+        isinstance(leaf, DTensor) and not isinstance(leaf, ShardTensor)
+        for leaf in torch.utils._pytree.tree_leaves((args, kwargs))
+    )
+
+
 class ShardTensor(torch.Tensor):
     r"""A distributed tensor class with support for uneven data sharding.
 
@@ -1629,10 +1637,12 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         elif str(func) in cls._named_function_registry and cls._enable_shard_patches:
             args, kwargs = _promote_plain_handler_args(args, kwargs)
             res = cls._named_function_registry[str(func)](func, types, args, kwargs)
-        elif _is_tracing(args, kwargs):
+        elif _is_tracing(args, kwargs) and not _has_plain_dtensor(args, kwargs):
             # Under torch.compile tracing, route unpatched ops straight to
             # ``__torch_dispatch__`` (like DTensor) so the aten-level path -- which AOT
             # handles correctly -- owns the op and keeps the graph differentiable.
+            # A plain DTensor operand (e.g. a Replicate parameter) would take that
+            # dispatch itself and reject the ShardTensor, so it uses the fallback.
             with torch._C.DisableTorchFunctionSubclass():
                 return func(*args, **kwargs)
         else:
diff --git a/test/domain_parallel/test_compile.py b/test/domain_parallel/test_compile.py
index ac37f8bbd3..771ebc227a 100644
--- a/test/domain_parallel/test_compile.py
+++ b/test/domain_parallel/test_compile.py
@@ -984,3 +984,36 @@ def f(x):
 @pytest.mark.timeout(300)
 def test_construct_shard_tensor_inside_compiled_fn_1d(distributed_mesh):
     run_construct_shard_tensor_inside_compiled_fn(distributed_mesh)
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.timeout(120)
+def test_compile_mixed_dtensor_shard_tensor_1d(distributed_mesh):
+    r"""A Replicate DTensor parameter meeting ShardTensors under compile (issue
+    #1986 #3): the DiT mask-token splice ``x * (1 - a) + token * a``.
+    Fwd + bwd match eager, including the DTensor parameter's gradient."""
+    from torch.distributed.tensor import distribute_tensor
+
+    dm = DistributedManager()
+    n = 4 * distributed_mesh.size(0) + 1  # uneven
+    torch.manual_seed(3)
+    x = torch.randn(1, n, 16, device=dm.device)
+    a = (torch.rand(1, n, 1, device=dm.device) > 0.5).float()
+
+    def run(compile_fn):
+        token = torch.nn.Parameter(
+            distribute_tensor(
+                torch.ones(1, 1, 16, device=dm.device), distributed_mesh, [Replicate()]
+            )
+        )
+        xs = scatter_tensor(x, 0, distributed_mesh, (Shard(1),), requires_grad=True)
+        as_ = scatter_tensor(a, 0, distributed_mesh, (Shard(1),))
+        out = compile_fn(lambda x, a: x * (1.0 - a) + token * a)(xs, as_)
+        out.square().sum().backward()
+        return out.full_tensor(), xs.grad.full_tensor(), token.grad.full_tensor()
+
+    torch._dynamo.reset()
+    eager = run(lambda f: f)
+    compiled = run(lambda f: torch.compile(f, backend="aot_eager", dynamic=False))
+    for e, c in zip(eager, compiled):
+        torch.testing.assert_close(c, e)

From 205896a3bf98d8e066bbd798f2a7b8dbe3712953 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Wed, 21 Oct 2026 13:37:12 -0700
Subject: [PATCH 5/8] ShardTensor: reduce plain-parameter grads under compile

Under tracing, ops with a grad-requiring plain operand (exact Tensor,
Parameter, FakeTensor or FunctionalTensor; grad mode on) take the DTensor
fallback, whose differentiable promotion reduces the Partial grad to a
plain tensor as in eager. Other tensor subclasses keep their dispatch.
---
 physicsnemo/domain_parallel/shard_tensor.py   | 45 +++++++++++---
 test/domain_parallel/test_compile.py          | 60 +++++++++++++++++++
 .../domain_parallel/test_tracing_promotion.py | 51 ++++++++++++++++
 3 files changed, 147 insertions(+), 9 deletions(-)
 create mode 100644 test/domain_parallel/test_tracing_promotion.py

diff --git a/physicsnemo/domain_parallel/shard_tensor.py b/physicsnemo/domain_parallel/shard_tensor.py
index 1c528fbdab..ae5b4990fd 100644
--- a/physicsnemo/domain_parallel/shard_tensor.py
+++ b/physicsnemo/domain_parallel/shard_tensor.py
@@ -27,7 +27,8 @@
 
 import torch
 import torch.distributed as dist
-from torch._subclasses.fake_tensor import is_fake
+from torch._subclasses.fake_tensor import FakeTensor, is_fake
+from torch._subclasses.functional_tensor import FunctionalTensor
 from torch.distributed.device_mesh import DeviceMesh, _mesh_resources
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._dtensor_spec import (
@@ -797,12 +798,36 @@ def _is_tracing(args: object, kwargs: object = None) -> bool:
     return False
 
 
-def _has_plain_dtensor(args: object, kwargs: object = None) -> bool:
-    r"""True when any operand is a ``DTensor`` that is not a ``ShardTensor``."""
-    return any(
-        isinstance(leaf, DTensor) and not isinstance(leaf, ShardTensor)
-        for leaf in torch.utils._pytree.tree_leaves((args, kwargs))
+# Plain tensors as they appear inside a traced ``__torch_function__``: a plain
+# tensor or parameter is a FakeTensor (FunctionalTensor during AOT's trace).
+# Exact types only: other subclasses (torchao, AsyncCollectiveTensor, ...) keep
+# their own type under tracing and must not be wrapped in a DTensor.
+_PLAIN_TRACED_TYPES = (
+    torch.Tensor,
+    torch.nn.Parameter,
+    FakeTensor,
+    FunctionalTensor,
+)
+
+
+def _needs_fallback_under_tracing(args: object, kwargs: object = None) -> bool:
+    r"""Whether a traced op must take the DTensor fallback instead of dispatch.
+
+    True for a plain ``DTensor`` operand (dispatch would reject the
+    ShardTensor) or a grad-requiring plain tensor to promote: the fallback's
+    differentiable promotion reduces its ``Partial`` gradient to a plain one,
+    as in eager; promotion at dispatch level records no backward.
+    """
+    promote = (
+        ShardTensor._promotion_mode is not TensorPromotionMode.DISABLED
+        and torch.is_grad_enabled()
     )
+    for leaf in torch.utils._pytree.tree_leaves((args, kwargs)):
+        if isinstance(leaf, DTensor) and not isinstance(leaf, ShardTensor):
+            return True
+        if promote and type(leaf) in _PLAIN_TRACED_TYPES and leaf.requires_grad:
+            return True
+    return False
 
 
 class ShardTensor(torch.Tensor):
@@ -1637,12 +1662,14 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         elif str(func) in cls._named_function_registry and cls._enable_shard_patches:
             args, kwargs = _promote_plain_handler_args(args, kwargs)
             res = cls._named_function_registry[str(func)](func, types, args, kwargs)
-        elif _is_tracing(args, kwargs) and not _has_plain_dtensor(args, kwargs):
+        elif _is_tracing(args, kwargs) and not _needs_fallback_under_tracing(
+            args, kwargs
+        ):
             # Under torch.compile tracing, route unpatched ops straight to
             # ``__torch_dispatch__`` (like DTensor) so the aten-level path -- which AOT
             # handles correctly -- owns the op and keeps the graph differentiable.
-            # A plain DTensor operand (e.g. a Replicate parameter) would take that
-            # dispatch itself and reject the ShardTensor, so it uses the fallback.
+            # Plain DTensor operands and grad-requiring plain tensors use the
+            # fallback instead (see ``_needs_fallback_under_tracing``).
             with torch._C.DisableTorchFunctionSubclass():
                 return func(*args, **kwargs)
         else:
diff --git a/test/domain_parallel/test_compile.py b/test/domain_parallel/test_compile.py
index 771ebc227a..0b00c02640 100644
--- a/test/domain_parallel/test_compile.py
+++ b/test/domain_parallel/test_compile.py
@@ -1010,6 +1010,8 @@ def run(compile_fn):
         as_ = scatter_tensor(a, 0, distributed_mesh, (Shard(1),))
         out = compile_fn(lambda x, a: x * (1.0 - a) + token * a)(xs, as_)
         out.square().sum().backward()
+        # FSDP2 requires a DTensor parameter's grad to be a plain DTensor.
+        assert type(token.grad) is DTensor, type(token.grad)
         return out.full_tensor(), xs.grad.full_tensor(), token.grad.full_tensor()
 
     torch._dynamo.reset()
@@ -1017,3 +1019,61 @@ def run(compile_fn):
     compiled = run(lambda f: torch.compile(f, backend="aot_eager", dynamic=False))
     for e, c in zip(eager, compiled):
         torch.testing.assert_close(c, e)
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.timeout(120)
+@pytest.mark.parametrize("compiled", [False, True])
+def test_shard_dtensor_param_grad_type_1d(distributed_mesh, compiled):
+    r"""A ``Shard(1)`` DTensor parameter (the recipe's pre-sharded ``pos_embed``)
+    added to ``Shard(1)`` tokens: its grad must be a plain DTensor, which FSDP2
+    requires for the reduce-scatter."""
+    from torch.distributed.tensor import distribute_tensor
+
+    dm = DistributedManager()
+    n = 4 * distributed_mesh.size(0)  # even, as for a static pos_embed
+    torch.manual_seed(4)
+    x = torch.randn(1, n, 16, device=dm.device)
+    pos = torch.randn(1, n, 16, device=dm.device)
+
+    pos_embed = torch.nn.Parameter(distribute_tensor(pos, distributed_mesh, [Shard(1)]))
+    xs = scatter_tensor(x, 0, distributed_mesh, (Shard(1),), requires_grad=True)
+
+    torch._dynamo.reset()
+    fn = lambda t: (t + pos_embed).square().sum()  # noqa: E731
+    if compiled:
+        fn = torch.compile(fn, backend="aot_eager", dynamic=False)
+    fn(xs).backward()
+
+    assert type(pos_embed.grad) is DTensor, type(pos_embed.grad)
+    assert pos_embed.grad.placements == (Shard(1),)
+    torch.testing.assert_close(pos_embed.grad.full_tensor(), 2 * (x + pos))
+
+
+@pytest.mark.multigpu_static
+@pytest.mark.timeout(120)
+@pytest.mark.parametrize("compiled", [False, True])
+def test_plain_param_grad_reduced_1d(distributed_mesh, compiled):
+    r"""A plain parameter meeting ``Shard(1)`` tokens gets a plain, fully
+    reduced gradient (the domain sum), eager and compiled. Under compile it
+    used to keep an unreduced ``ShardTensor(Partial)`` grad, which FSDP2
+    rejects (the DiT's modulation weights and mask tokens)."""
+    dm = DistributedManager()
+    n = 4 * distributed_mesh.size(0) + 1  # uneven
+    torch.manual_seed(5)
+    x = torch.randn(1, n, 16, device=dm.device)
+    w = torch.randn(1, 1, 16, device=dm.device)
+
+    param = torch.nn.Parameter(w.clone())
+    xs = scatter_tensor(x, 0, distributed_mesh, (Shard(1),), requires_grad=True)
+
+    torch._dynamo.reset()
+    fn = lambda t: (t * param).square().sum()  # noqa: E731
+    if compiled:
+        fn = torch.compile(fn, backend="aot_eager", dynamic=False)
+    fn(xs).backward()
+
+    ref = w.clone().requires_grad_(True)
+    (x * ref).square().sum().backward()
+    assert type(param.grad) is torch.Tensor, type(param.grad)
+    torch.testing.assert_close(param.grad, ref.grad)
diff --git a/test/domain_parallel/test_tracing_promotion.py b/test/domain_parallel/test_tracing_promotion.py
new file mode 100644
index 0000000000..7eb218e211
--- /dev/null
+++ b/test/domain_parallel/test_tracing_promotion.py
@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""Which operands send a traced ShardTensor op to the DTensor fallback.
+
+Only plain tensors (as they appear under tracing) that need a gradient are
+promoted there; other tensor subclasses keep their own dispatch.
+"""
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.testing._internal.two_tensor import TwoTensor
+
+from physicsnemo.domain_parallel.shard_tensor import _needs_fallback_under_tracing
+
+
+def test_fake_plain_parameter_needs_fallback():
+    with FakeTensorMode():
+        param = torch.nn.Parameter(torch.randn(4))
+    assert _needs_fallback_under_tracing((param,))
+
+
+def test_plain_without_grad_takes_dispatch():
+    with FakeTensorMode():
+        buf = torch.randn(4)
+    assert not _needs_fallback_under_tracing((buf,))
+
+
+def test_other_subclass_takes_dispatch():
+    t = TwoTensor(torch.randn(4), torch.randn(4)).requires_grad_(True)
+    assert not _needs_fallback_under_tracing((t,))
+
+
+def test_no_grad_mode_takes_dispatch():
+    with FakeTensorMode():
+        param = torch.nn.Parameter(torch.randn(4))
+    with torch.no_grad():
+        assert not _needs_fallback_under_tracing((param,))

From 6d90953ff9765e5692582ba9ac44869bc4a38ff2 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Wed, 28 Oct 2026 10:45:23 -0700
Subject: [PATCH 6/8] CHANGELOG: ShardTensor compile fixes

---
 CHANGELOG.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc4e708b8d..163df3ccfe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -394,6 +394,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Padé networks (pytorch/pytorch#193194).
 - Preserve unreduced ShardTensor placement dimensions in `sum` and `mean`
   when `keepdim=True`, including independently sharded simulation and spatial axes.
+- `torch.compile` with ShardTensor: neighborhood attention (`na1d`/`na2d`/
+  `na3d`) keeps its halo exchange when compiled, `select` handles a negative
+  `dim`, dynamic shapes no longer fail in the spec hash, a plain DTensor
+  operand next to a ShardTensor uses the DTensor fallback, and a plain
+  parameter used with a ShardTensor gets a reduced gradient as in eager.
 - Multi-node jobs launched with `srun` no longer hang in
   `DistributedManager.initialize()` when Slurm starts rank 0 on a different node
   from the one that ran `srun`. The rendezvous now uses the job step's first node.

From 57af780dac20d7f5a4045fcc474cca9a895d7e89 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:26:19 -0700
Subject: [PATCH 7/8] ShardTensorSpec: don't compare SymInts across traces in
 __eq__

Symbolic specs hash by (mesh, placements) only, so specs from different
compiles collide in the spec cache. Comparing their SymInts raised
'SymInt ... not tracked' (1-GPU UNet force_sharding compile). Compare
ints by value and SymInts by identity; anything else is a cache miss.
---
 .../domain_parallel/_shard_tensor_spec.py     | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/physicsnemo/domain_parallel/_shard_tensor_spec.py b/physicsnemo/domain_parallel/_shard_tensor_spec.py
index e5b5e23a52..934f486e4b 100644
--- a/physicsnemo/domain_parallel/_shard_tensor_spec.py
+++ b/physicsnemo/domain_parallel/_shard_tensor_spec.py
@@ -34,6 +34,21 @@
 from physicsnemo.distributed.utils import compute_split_shapes
 
 
+def _is_symbolic(spec: "ShardTensorSpec") -> bool:
+    r"""Whether the spec's global shape or stride holds a SymInt."""
+    meta = spec.tensor_meta
+    return meta is not None and any(
+        type(d) is not int for d in (*meta.shape, *meta.stride)
+    )
+
+
+def _same_dims(a, b) -> bool:
+    r"""Equal dims without SymInt arithmetic: ints by value, SymInts by identity."""
+    return len(a) == len(b) and all(
+        x is y or (type(x) is int and type(y) is int and x == y) for x, y in zip(a, b)
+    )
+
+
 @dataclass(kw_only=True)
 class ShardTensorSpec(DTensorSpec):
     r"""A distributed tensor specification that tracks sharding information.
@@ -222,6 +237,34 @@ def __eq__(self, other: object) -> bool:
         """
         if not isinstance(other, ShardTensorSpec):
             return False
+        if _is_symbolic(self) or _is_symbolic(other):
+            # Symbolic shapes hash by (mesh, placements) only, so collisions
+            # with specs from other traces reach here: never compare SymInts
+            # (that records guards or fails across traces). Same objects or
+            # equal ints only; anything else is a cache miss, not an error.
+            if self is other:
+                return True
+            a, b = self.tensor_meta, other.tensor_meta
+            if (
+                self.mesh != other.mesh
+                or self.placements != other.placements
+                or (a is None) != (b is None)
+            ):
+                return False
+            if a is not None and not (
+                a.dtype == b.dtype
+                and _same_dims(a.shape, b.shape)
+                and _same_dims(a.stride, b.stride)
+            ):
+                return False
+            sa, sb = self._sharding_shapes, other._sharding_shapes
+            if sa is None or sb is None:
+                return sa is sb
+            return sa.keys() == sb.keys() and all(
+                len(sa[k]) == len(sb[k])
+                and all(_same_dims(x, y) for x, y in zip(sa[k], sb[k]))
+                for k in sa
+            )
         if not super().__eq__(other):
             return False
         if self._sharding_shapes != other._sharding_shapes:

From 93dc28f423d9cf9caaf75c16adf0c4af94365a49 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:26:20 -0700
Subject: [PATCH 8/8] NATTEN: route subclass inputs through
 physicsnemo::na{1,2,3}d ops

Dynamo treats ShardTensor as a plain tensor and never calls its
__torch_function__, so a Python-level has_torch_function check (or
allow_in_graph) can't reach the ShardTensor handler cleanly. na1d/2d/3d
now call torch.library custom ops when an operand is a tensor subclass;
ShardTensor registers the halo-exchange handler on them with
register_named_function_handler, which runs in eager and under AOT.
Plain tensors still call NATTEN directly.
---
 .../shard_utils/natten_patches.py             |  76 ++++----
 physicsnemo/nn/functional/natten.py           | 166 +++++++++---------
 .../ops/test_natten_compile.py                |   9 +-
 test/nn/functional/test_natten.py             |   8 +-
 4 files changed, 124 insertions(+), 135 deletions(-)

diff --git a/physicsnemo/domain_parallel/shard_utils/natten_patches.py b/physicsnemo/domain_parallel/shard_utils/natten_patches.py
index bce4dce342..5b72019f44 100644
--- a/physicsnemo/domain_parallel/shard_utils/natten_patches.py
+++ b/physicsnemo/domain_parallel/shard_utils/natten_patches.py
@@ -18,7 +18,6 @@
 
 from typing import Any, Callable
 
-import torch
 from torch.distributed.tensor.placement_types import Shard
 
 from physicsnemo.core.version_check import OptionalImport
@@ -33,14 +32,11 @@
     MissingShardPatch,
     UndeterminedShardingError,
 )
-from physicsnemo.nn.functional.natten import na1d, na2d, na3d
+from physicsnemo.nn.functional import natten as _na_ops  # noqa: F401 (defines the ops)
 
 _natten = OptionalImport("natten")
-_raw_func_map = {
-    na1d: lambda: _natten.functional.na1d,
-    na2d: lambda: _natten.functional.na2d,
-    na3d: lambda: _natten.functional.na3d,
-}
+# Argument order of the ``physicsnemo::na{1,2,3}d`` ops.
+_OP_ARGS = ("q", "k", "v", "kernel_size", "dilation", "is_causal", "scale", "backend")
 
 __all__ = ["na1d_wrapper", "na2d_wrapper", "na3d_wrapper"]
 
@@ -226,54 +222,55 @@ def _natten_wrapper(
     types: tuple[Any, ...],
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
-) -> torch.Tensor | ShardTensor:
-    r"""Shared wrapper for natten functions to support sharded tensors.
+) -> ShardTensor:
+    r"""ShardTensor handler for the ``physicsnemo::na{1,2,3}d`` ops.
 
-    Registered with :meth:`ShardTensor.register_function_handler` so that calls
-    to :func:`~physicsnemo.nn.functional.natten.na1d`,
-    :func:`~physicsnemo.nn.functional.natten.na2d`, or
-    :func:`~physicsnemo.nn.functional.natten.na3d` automatically route through
-    this handler when any argument is a :class:`ShardTensor`.
+    :func:`~physicsnemo.nn.functional.natten.na1d`,
+    :func:`~physicsnemo.nn.functional.natten.na2d` and
+    :func:`~physicsnemo.nn.functional.natten.na3d` call these ops for tensor
+    subclass operands. Registered by op name, so ``__torch_function__`` routes
+    the op here in eager and under ``torch.compile``; the halo exchange and
+    the local NATTEN call carry autograd.
 
     Parameters
     ----------
     func : Callable
-        The wrapped natten function (passed by ``__torch_function__``).
+        The op overload, e.g. ``torch.ops.physicsnemo.na2d.default``.
     types : tuple[Any, ...]
         The types of the inputs (unused).
     args : tuple[Any, ...]
-        Positional arguments containing query, key, value tensors and kernel_size.
+        Op arguments, in the order of ``_OP_ARGS``.
     kwargs : dict[str, Any]
-        Keyword arguments including ``dilation``.
+        Op arguments passed by keyword.
 
     Returns
     -------
-    Union[torch.Tensor, ShardTensor]
-        Result tensor as either ``torch.Tensor`` or ShardTensor depending on input types.
+    ShardTensor
+        Neighborhood attention output, sharded like ``q``.
 
     Raises
     ------
     UndeterminedShardingError
-        If input tensor types are mismatched.
+        If q, k and v are not all ShardTensors.
     """
-    q, k, v, kernel_size = args[0], args[1], args[2], args[3]
-
-    dilation = kwargs.get("dilation", 1)
-    natten_kwargs = {_k: _v for _k, _v in kwargs.items() if _k != "dilation"}
-
-    if all(type(_t) is torch.Tensor for _t in (q, k, v)):
-        return func(
-            q, k, v, kernel_size=kernel_size, dilation=dilation, **natten_kwargs
-        )
-    elif all(isinstance(_t, ShardTensor) for _t in (q, k, v)):
-        raw_func = _raw_func_map[func]()
-        return _partial_natten(
-            q, k, v, kernel_size, dilation, base_func=raw_func, **natten_kwargs
-        )
-    else:
+    bound = {**dict(zip(_OP_ARGS, args)), **kwargs}
+    q, k, v = bound["q"], bound["k"], bound["v"]
+    if not all(isinstance(_t, ShardTensor) for _t in (q, k, v)):
         raise UndeterminedShardingError(
-            "q, k, and v must all be the same types (torch.Tensor or ShardTensor)"
+            "q, k, and v must all be ShardTensors for sharded neighborhood attention"
         )
+    raw_func = getattr(_natten.functional, func.__name__.split(".")[0])
+    return _partial_natten(
+        q,
+        k,
+        v,
+        bound["kernel_size"],
+        bound["dilation"],
+        base_func=raw_func,
+        is_causal=bound["is_causal"],
+        scale=bound["scale"],
+        backend=bound["backend"],
+    )
 
 
 # Public aliases for explicit registration
@@ -281,6 +278,7 @@ def _natten_wrapper(
 na2d_wrapper = _natten_wrapper
 na3d_wrapper = _natten_wrapper
 
-ShardTensor.register_function_handler(na1d, na1d_wrapper)
-ShardTensor.register_function_handler(na2d, na2d_wrapper)
-ShardTensor.register_function_handler(na3d, na3d_wrapper)
+for _name in ("na1d", "na2d", "na3d"):
+    ShardTensor.register_named_function_handler(
+        f"physicsnemo.{_name}.default", _natten_wrapper
+    )
diff --git a/physicsnemo/nn/functional/natten.py b/physicsnemo/nn/functional/natten.py
index 69ad827730..46be81fbf6 100644
--- a/physicsnemo/nn/functional/natten.py
+++ b/physicsnemo/nn/functional/natten.py
@@ -16,47 +16,80 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Optional
 
 import torch
-from torch.overrides import handle_torch_function, has_torch_function
 
 from physicsnemo.core.version_check import OptionalImport
 
 _natten = OptionalImport("natten")
 
+# Options the ``physicsnemo::na{1,2,3}d`` ops carry, with NATTEN's defaults.
+_OP_OPTIONS = {"is_causal": False, "scale": None, "backend": None}
 
-def _subclass_under_compile(*tensors: torch.Tensor) -> bool:
-    r"""True while Dynamo traces a call with a tensor-subclass operand.
 
-    Dynamo folds ``has_torch_function`` to ``False`` for traceable wrapper
-    subclasses such as ``ShardTensor``, so the check above alone would hand the
-    subclass to the raw kernel. ``type()`` still reports the subclass.
+def _register_na_op(name: str) -> torch._ops.OpOverload:
+    r"""Define ``physicsnemo::<name>``: NATTEN's ``<name>`` as a graph op.
+
+    Tensor subclasses such as ``ShardTensor`` intercept the op by name in
+    ``__torch_function__`` (``register_named_function_handler``), in eager
+    and under ``torch.compile`` alike, and handle autograd themselves.
     """
-    return torch.compiler.is_compiling() and any(
-        type(t) is not torch.Tensor for t in tensors
-    )
 
+    @torch.library.custom_op(f"physicsnemo::{name}", mutates_args=())
+    def op(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        kernel_size: int,
+        dilation: int,
+        is_causal: bool,
+        scale: Optional[float],
+        backend: Optional[str],
+    ) -> torch.Tensor:
+        return getattr(_natten.functional, name)(
+            q,
+            k,
+            v,
+            kernel_size,
+            dilation=dilation,
+            is_causal=is_causal,
+            scale=scale,
+            backend=backend,
+        )
 
-@torch._dynamo.allow_in_graph
-def _dispatch_in_graph(
-    name: str,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    kernel_size: int,
-    dilation: int,
-    kwargs: dict[str, Any],
-) -> torch.Tensor:
-    r"""Subclass dispatch as one opaque graph node.
+    @op.register_fake
+    def _(q, k, v, kernel_size, dilation, is_causal, scale, backend):
+        return torch.empty_like(q)
 
-    Dynamo records the call and runs it under fake propagation, where
-    ``__torch_function__`` dispatches to the subclass handler (for
-    ``ShardTensor``, the halo-exchange path), as for ``torch`` ops.
-    """
-    func = _NA_FUNCS[name]
-    return handle_torch_function(
-        func, (q, k, v), q, k, v, kernel_size, dilation=dilation, **kwargs
+    return op._opoverload
+
+
+_NA_OPS = {name: _register_na_op(name) for name in ("na1d", "na2d", "na3d")}
+
+
+def _na(name, q, k, v, kernel_size, dilation, kwargs):
+    r"""Call NATTEN on plain tensors; route tensor subclasses through the op."""
+    if all(type(t) in (torch.Tensor, torch.nn.Parameter) for t in (q, k, v)):
+        return getattr(_natten.functional, name)(
+            q, k, v, kernel_size, dilation=dilation, **kwargs
+        )
+    unsupported = set(kwargs) - set(_OP_OPTIONS)
+    if unsupported:
+        raise TypeError(
+            f"{name} on tensor subclasses supports only {sorted(_OP_OPTIONS)}, "
+            f"got {sorted(unsupported)}"
+        )
+    opts = {**_OP_OPTIONS, **kwargs}
+    return _NA_OPS[name](
+        q,
+        k,
+        v,
+        kernel_size,
+        dilation,
+        opts["is_causal"],
+        opts["scale"],
+        opts["backend"],
     )
 
 
@@ -68,12 +101,12 @@ def na1d(
     dilation: int = 1,
     **kwargs: Any,
 ) -> torch.Tensor:
-    r"""Compute 1D neighborhood attention, with ``__torch_function__`` dispatch.
+    r"""Compute 1D neighborhood attention, with tensor-subclass dispatch.
 
-    This is a thin wrapper around :func:`natten.functional.na1d` that enables
-    automatic dispatch through PyTorch's ``__torch_function__`` protocol. When
-    called with a tensor subclass (e.g. ``ShardTensor``), the registered handler
-    is invoked instead of the underlying natten implementation.
+    This is a thin wrapper around :func:`natten.functional.na1d`. Tensor
+    subclass operands (e.g. ``ShardTensor``) go through the
+    ``physicsnemo::na1d`` op, which subclasses intercept by name; there only
+    ``is_causal``, ``scale`` and ``backend`` are accepted as extra options.
 
     Parameters
     ----------
@@ -96,20 +129,7 @@ def na1d(
     torch.Tensor
         Output tensor of the same shape as ``q``.
     """
-    if has_torch_function((q, k, v)):
-        return handle_torch_function(
-            na1d,
-            (q, k, v),
-            q,
-            k,
-            v,
-            kernel_size,
-            dilation=dilation,
-            **kwargs,
-        )
-    if _subclass_under_compile(q, k, v):
-        return _dispatch_in_graph("na1d", q, k, v, kernel_size, dilation, kwargs)
-    return _natten.functional.na1d(q, k, v, kernel_size, dilation=dilation, **kwargs)
+    return _na("na1d", q, k, v, kernel_size, dilation, kwargs)
 
 
 def na2d(
@@ -120,12 +140,12 @@ def na2d(
     dilation: int = 1,
     **kwargs: Any,
 ) -> torch.Tensor:
-    r"""Compute 2D neighborhood attention, with ``__torch_function__`` dispatch.
+    r"""Compute 2D neighborhood attention, with tensor-subclass dispatch.
 
-    This is a thin wrapper around :func:`natten.functional.na2d` that enables
-    automatic dispatch through PyTorch's ``__torch_function__`` protocol. When
-    called with a tensor subclass (e.g. ``ShardTensor``), the registered handler
-    is invoked instead of the underlying natten implementation.
+    This is a thin wrapper around :func:`natten.functional.na2d`. Tensor
+    subclass operands (e.g. ``ShardTensor``) go through the
+    ``physicsnemo::na2d`` op, which subclasses intercept by name; there only
+    ``is_causal``, ``scale`` and ``backend`` are accepted as extra options.
 
     Parameters
     ----------
@@ -148,20 +168,7 @@ def na2d(
     torch.Tensor
         Output tensor of the same shape as ``q``.
     """
-    if has_torch_function((q, k, v)):
-        return handle_torch_function(
-            na2d,
-            (q, k, v),
-            q,
-            k,
-            v,
-            kernel_size,
-            dilation=dilation,
-            **kwargs,
-        )
-    if _subclass_under_compile(q, k, v):
-        return _dispatch_in_graph("na2d", q, k, v, kernel_size, dilation, kwargs)
-    return _natten.functional.na2d(q, k, v, kernel_size, dilation=dilation, **kwargs)
+    return _na("na2d", q, k, v, kernel_size, dilation, kwargs)
 
 
 def na3d(
@@ -172,12 +179,12 @@ def na3d(
     dilation: int = 1,
     **kwargs: Any,
 ) -> torch.Tensor:
-    r"""Compute 3D neighborhood attention, with ``__torch_function__`` dispatch.
+    r"""Compute 3D neighborhood attention, with tensor-subclass dispatch.
 
-    This is a thin wrapper around :func:`natten.functional.na3d` that enables
-    automatic dispatch through PyTorch's ``__torch_function__`` protocol. When
-    called with a tensor subclass (e.g. ``ShardTensor``), the registered handler
-    is invoked instead of the underlying natten implementation.
+    This is a thin wrapper around :func:`natten.functional.na3d`. Tensor
+    subclass operands (e.g. ``ShardTensor``) go through the
+    ``physicsnemo::na3d`` op, which subclasses intercept by name; there only
+    ``is_causal``, ``scale`` and ``backend`` are accepted as extra options.
 
     Parameters
     ----------
@@ -200,22 +207,7 @@ def na3d(
     torch.Tensor
         Output tensor of the same shape as ``q``.
     """
-    if has_torch_function((q, k, v)):
-        return handle_torch_function(
-            na3d,
-            (q, k, v),
-            q,
-            k,
-            v,
-            kernel_size,
-            dilation=dilation,
-            **kwargs,
-        )
-    if _subclass_under_compile(q, k, v):
-        return _dispatch_in_graph("na3d", q, k, v, kernel_size, dilation, kwargs)
-    return _natten.functional.na3d(q, k, v, kernel_size, dilation=dilation, **kwargs)
-
+    return _na("na3d", q, k, v, kernel_size, dilation, kwargs)
 
-_NA_FUNCS = {"na1d": na1d, "na2d": na2d, "na3d": na3d}
 
 __all__ = ["na1d", "na2d", "na3d"]
diff --git a/test/domain_parallel/ops/test_natten_compile.py b/test/domain_parallel/ops/test_natten_compile.py
index a3d046d1d2..263029dda3 100644
--- a/test/domain_parallel/ops/test_natten_compile.py
+++ b/test/domain_parallel/ops/test_natten_compile.py
@@ -16,11 +16,10 @@
 
 r"""``torch.compile`` over sharded neighborhood attention (issue #1986).
 
-Dynamo treats ``ShardTensor`` as a traceable wrapper subclass and folds
-``has_torch_function`` to ``False`` for it, so a traced ``na2d`` would call the
-raw kernel on the ShardTensor instead of the halo-exchanging handler. The raw
-kernel here is a pure-torch reference that, like an opaque kernel, refuses
-distributed inputs; it also lets the test run without NATTEN installed.
+A traced ``na2d`` on ShardTensors must emit the ``physicsnemo::na2d`` op so
+ShardTensor's named handler (halo exchange) runs, rather than the raw kernel
+seeing the ShardTensor. The raw kernel here is a pure-torch reference that,
+like an opaque kernel, refuses distributed inputs; it also lets the test run without NATTEN installed.
 """
 
 import pytest
diff --git a/test/nn/functional/test_natten.py b/test/nn/functional/test_natten.py
index 90d8e48574..140f0401ca 100644
--- a/test/nn/functional/test_natten.py
+++ b/test/nn/functional/test_natten.py
@@ -20,7 +20,7 @@
 
 * produce bit-identical results to the underlying ``natten.functional`` calls,
 * preserve tensor shapes through forward and backward,
-* correctly dispatch through ``__torch_function__`` for tensor subclasses,
+* route tensor subclasses through the ``physicsnemo::na{1,2,3}d`` ops,
 * degenerate to standard scaled dot-product attention when the kernel window
   covers the entire spatial extent.
 """
@@ -148,7 +148,7 @@ def test_torch_function_dispatch(self, device):
         _DispatchRecorder.dispatched_funcs.clear()
         na1d(q, k, v, kernel_size=3)
 
-        assert na1d in _DispatchRecorder.dispatched_funcs
+        assert torch.ops.physicsnemo.na1d.default in _DispatchRecorder.dispatched_funcs
 
     def test_full_window_matches_sdpa(self, device):
         """When kernel covers the entire sequence, NA degenerates to SDPA."""
@@ -242,7 +242,7 @@ def test_torch_function_dispatch(self, device):
         _DispatchRecorder.dispatched_funcs.clear()
         na2d(q, k, v, kernel_size=3)
 
-        assert na2d in _DispatchRecorder.dispatched_funcs
+        assert torch.ops.physicsnemo.na2d.default in _DispatchRecorder.dispatched_funcs
 
     def test_full_window_matches_sdpa(self, device):
         """When kernel covers the full spatial extent, NA degenerates to SDPA."""
@@ -334,7 +334,7 @@ def test_torch_function_dispatch(self, device):
         _DispatchRecorder.dispatched_funcs.clear()
         na3d(q, k, v, kernel_size=3)
 
-        assert na3d in _DispatchRecorder.dispatched_funcs
+        assert torch.ops.physicsnemo.na3d.default in _DispatchRecorder.dispatched_funcs
 
     def test_full_window_matches_sdpa(self, device):
         """When kernel covers the full spatial extent, NA degenerates to SDPA."""