From 075c5c56769a03731b427e341fd4d629faca12f3 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Mon, 13 Jul 2026 12:07:11 -0600
Subject: [PATCH 01/10] Add model-agnostic shrink-and-perturb weight
 re-initialization to physicsnemo.nn

---
 CHANGELOG.md                               |   6 +
 docs/api/nn/layers/weight_init.rst         |   6 +
 docs/api/physicsnemo.nn.layers.rst         |   1 +
 physicsnemo/nn/__init__.py                 |   1 +
 physicsnemo/nn/module/utils/__init__.py    |   2 +-
 physicsnemo/nn/module/utils/weight_init.py | 120 ++++++++++++++++++
 test/nn/module/test_weight_init.py         | 135 +++++++++++++++++++++
 7 files changed, 270 insertions(+), 1 deletion(-)
 create mode 100644 docs/api/nn/layers/weight_init.rst
 create mode 100644 test/nn/module/test_weight_init.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2fd9e860ad..4f04773beb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Adds `physicsnemo.nn.shrink_and_perturb_`, an in-place shrink-and-perturb
+  weight re-initialization (Ash & Adams, NeurIPS 2020) for warm-starting from
+  pretrained weights. Model-agnostic, with a per-tensor scale-aware Gaussian
+  default (`theta <- shrink*theta + perturb*std(theta)*noise`), a plain-Gaussian
+  and custom-callable noise source, and an `include` filter to restrict the
+  perturbation to (e.g.) transferred backbone parameters.
 - Adds a `global_shape` argument to `ShardTensor.from_local`, enabling the
   no-communication `sharding_shapes="chunk"` path.
 - Adds exact-boundary quality mesh generation to
diff --git a/docs/api/nn/layers/weight_init.rst b/docs/api/nn/layers/weight_init.rst
new file mode 100644
index 0000000000..a387191198
--- /dev/null
+++ b/docs/api/nn/layers/weight_init.rst
@@ -0,0 +1,6 @@
+Weight Initialization
+=====================
+
+.. automodule:: physicsnemo.nn.module.utils.weight_init
+   :members:
+   :show-inheritance:
diff --git a/docs/api/physicsnemo.nn.layers.rst b/docs/api/physicsnemo.nn.layers.rst
index e472934730..04122df807 100644
--- a/docs/api/physicsnemo.nn.layers.rst
+++ b/docs/api/physicsnemo.nn.layers.rst
@@ -18,3 +18,4 @@ PhysicsNeMo Layers
    nn/layers/regularization
    nn/layers/specialized
    nn/layers/graph_geometry
+   nn/layers/weight_init
diff --git a/physicsnemo/nn/__init__.py b/physicsnemo/nn/__init__.py
index 4a70fa4d3d..575a2b4fe4 100644
--- a/physicsnemo/nn/__init__.py
+++ b/physicsnemo/nn/__init__.py
@@ -158,5 +158,6 @@
     SwinTransformer,
 )
 from .module.unet_layers import UNetBlock
+from .module.utils.weight_init import shrink_and_perturb_
 from .module.weight_fact import WeightFactLinear
 from .module.weight_norm import WeightNormLinear
diff --git a/physicsnemo/nn/module/utils/__init__.py b/physicsnemo/nn/module/utils/__init__.py
index 2881d37df8..ebcba88f0b 100644
--- a/physicsnemo/nn/module/utils/__init__.py
+++ b/physicsnemo/nn/module/utils/__init__.py
@@ -32,4 +32,4 @@
     get_pad2d,
     get_pad3d,
 )
-from .weight_init import trunc_normal_
+from .weight_init import shrink_and_perturb_, trunc_normal_
diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 10af9b7561..87910ddf98 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import warnings
+from typing import Callable, Optional, Union
 
 import numpy as np
 import torch
@@ -39,6 +40,125 @@ def trunc_normal_(*args, **kwargs):
     return _torch_trunc_normal_(*args, **kwargs)
 
 
+_NOISE_PRESETS = ("scaled_normal", "normal")
+
+
+@torch.no_grad()
+def shrink_and_perturb_(
+    module: torch.nn.Module,
+    shrink: float = 0.5,
+    perturb: float = 0.1,
+    *,
+    noise: Union[str, Callable[[str, torch.Tensor], torch.Tensor]] = "scaled_normal",
+    include: Optional[Callable[[str, torch.Tensor], bool]] = None,
+    generator: Optional[torch.Generator] = None,
+) -> torch.nn.Module:
+    r"""Apply *shrink-and-perturb* re-initialization to ``module`` in place.
+
+    For every selected parameter :math:`\theta`, the update is
+
+    .. math:: \theta \leftarrow \lambda\,\theta + p\,\varepsilon,
+
+    where :math:`\lambda` is ``shrink``, :math:`p` is ``perturb``, and
+    :math:`\varepsilon` is fresh noise. Shrinking a *pretrained* weight toward
+    zero restores the scale statistics and plasticity of a fresh initialization
+    while the noise breaks symmetry, yet the shrink term keeps the direction of
+    the pretrained features. In warm-started training this often reaches a lower
+    loss asymptote than fine-tuning the raw pretrained weights (Ash & Adams,
+    *On Warm-Starting Neural Network Training*, NeurIPS 2020).
+
+    The operation is only meaningful on **pretrained** weights: applied to a
+    fresh initialization it merely rescales and re-noises random values.
+
+    Parameters
+    ----------
+    module : torch.nn.Module
+        Module whose parameters are perturbed in place. Buffers (e.g.
+        batch-norm running statistics) are left untouched.
+    shrink : float, optional
+        Multiplicative retention factor :math:`\lambda \ge 0` applied to each
+        weight. Values in ``[0, 1)`` shrink toward zero; ``1.0`` disables the
+        shrink. Default ``0.5``.
+    perturb : float, optional
+        Noise scale :math:`p \ge 0`. With ``noise="scaled_normal"`` this is the
+        noise level relative to each tensor's own standard deviation. Default
+        ``0.1``.
+    noise : str or callable, optional
+        Source of the perturbation :math:`\varepsilon`:
+
+        - ``"scaled_normal"`` (default):
+          :math:`\varepsilon = \operatorname{std}(\theta)\,z` with
+          :math:`z \sim \mathcal{N}(0, 1)`, i.e. Gaussian noise scaled by the
+          per-tensor standard deviation of the pre-shrink weight. Scale aware
+          and free of any architectural assumptions (expects at least two
+          elements per parameter tensor).
+        - ``"normal"``: :math:`\varepsilon = z \sim \mathcal{N}(0, 1)`,
+          unscaled.
+        - a callable ``(name, param) -> tensor`` returning
+          :math:`\varepsilon`, e.g. ``lambda n, p: fresh[n]`` to interpolate
+          toward the weights of a freshly constructed reference model.
+    include : callable, optional
+        Predicate ``(name, param) -> bool`` selecting which parameters to
+        perturb. Parameters for which it returns ``False`` are left entirely
+        unchanged (not even shrunk). Default: all parameters. For warm-starting,
+        pass e.g. ``include=lambda n, p: n in transferred`` to perturb only the
+        transferred backbone.
+    generator : torch.Generator, optional
+        Generator for the built-in Gaussian noise, for reproducibility. Must be
+        on the same device as ``module``'s parameters. Ignored when ``noise``
+        is a callable.
+
+    Returns
+    -------
+    torch.nn.Module
+        The same ``module``, modified in place (returned for chaining).
+
+    Raises
+    ------
+    ValueError
+        If ``shrink`` or ``perturb`` is negative, or ``noise`` is an unknown
+        string.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from physicsnemo.nn import shrink_and_perturb_
+    >>> model = torch.nn.Linear(4, 4)
+    >>> _ = shrink_and_perturb_(model, shrink=0.6, perturb=0.1)
+    """
+    if shrink < 0.0 or perturb < 0.0:
+        raise ValueError(
+            f"shrink and perturb must be non-negative, got shrink={shrink}, "
+            f"perturb={perturb}"
+        )
+
+    if callable(noise):
+        noise_fn = noise
+    elif noise == "scaled_normal":
+
+        def noise_fn(name: str, p: torch.Tensor) -> torch.Tensor:
+            z = torch.empty_like(p).normal_(generator=generator)
+            return p.detach().std() * z
+
+    elif noise == "normal":
+
+        def noise_fn(name: str, p: torch.Tensor) -> torch.Tensor:
+            return torch.empty_like(p).normal_(generator=generator)
+
+    else:
+        raise ValueError(
+            f'Invalid noise "{noise}"; expected a callable or one of {_NOISE_PRESETS}'
+        )
+
+    for name, p in module.named_parameters():
+        if include is not None and not include(name, p):
+            continue
+        # eps is computed from the pre-shrink weight (matters for "scaled_normal").
+        eps = noise_fn(name, p)
+        p.mul_(shrink).add_(eps, alpha=perturb)
+    return module
+
+
 def _weight_init(shape: tuple, mode: str, fan_in: int, fan_out: int):
     """
     Unified routine for initializing weights and biases.
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
new file mode 100644
index 0000000000..9aad446e3b
--- /dev/null
+++ b/test/nn/module/test_weight_init.py
@@ -0,0 +1,135 @@
+# 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.
+
+import pytest
+import torch
+
+from physicsnemo.nn import shrink_and_perturb_
+
+
+def test_shrink_and_perturb_shrink_only(device):
+    """perturb=0 shrinks every param exactly; returns the same module."""
+    model = torch.nn.Linear(8, 8).to(device)
+    orig = {n: p.detach().clone() for n, p in model.named_parameters()}
+
+    ret = shrink_and_perturb_(model, shrink=0.4, perturb=0.0)
+
+    assert ret is model
+    for n, p in model.named_parameters():
+        assert torch.allclose(p, 0.4 * orig[n])
+
+
+def test_shrink_and_perturb_noop(device):
+    """shrink=1, perturb=0 leaves every param untouched."""
+    model = torch.nn.Linear(8, 8).to(device)
+    orig = {n: p.detach().clone() for n, p in model.named_parameters()}
+
+    shrink_and_perturb_(model, shrink=1.0, perturb=0.0)
+
+    for n, p in model.named_parameters():
+        assert torch.equal(p, orig[n])
+
+
+def test_shrink_and_perturb_callable_noise_exact(device):
+    """Callable noise reproduces theta <- shrink*theta + perturb*eps exactly."""
+    model = torch.nn.Linear(4, 4).to(device)
+    orig = {n: p.detach().clone() for n, p in model.named_parameters()}
+    eps = {n: torch.ones_like(p) for n, p in model.named_parameters()}
+
+    shrink_and_perturb_(model, shrink=0.5, perturb=0.2, noise=lambda n, p: eps[n])
+
+    for n, p in model.named_parameters():
+        assert torch.allclose(p, 0.5 * orig[n] + 0.2 * eps[n])
+
+
+def test_scaled_normal_magnitude(device):
+    """scaled_normal noise std tracks perturb * std(theta) per tensor."""
+    model = torch.nn.Linear(256, 256).to(device)
+    with torch.no_grad():
+        model.weight.normal_(0.0, 3.0)
+    w_std = model.weight.std().item()
+
+    gen = torch.Generator(device=device).manual_seed(0)
+    # shrink=0 isolates the noise term: result == perturb * std(theta) * z.
+    shrink_and_perturb_(
+        model,
+        shrink=0.0,
+        perturb=0.1,
+        noise="scaled_normal",
+        include=lambda n, p: n == "weight",
+        generator=gen,
+    )
+
+    expected = 0.1 * w_std
+    assert abs(model.weight.std().item() - expected) / expected < 0.1
+
+
+@pytest.mark.parametrize("noise", ["scaled_normal", "normal"])
+def test_gaussian_noise_reproducible(device, noise):
+    """Same generator seed -> identical result; different seed -> different."""
+
+    def run(seed):
+        torch.manual_seed(0)  # identical init across runs
+        model = torch.nn.Linear(8, 8).to(device)
+        gen = torch.Generator(device=device).manual_seed(seed)
+        shrink_and_perturb_(model, noise=noise, generator=gen)
+        return {n: p.detach().clone() for n, p in model.named_parameters()}
+
+    a, b, c = run(123), run(123), run(999)
+    for n in a:
+        assert torch.equal(a[n], b[n])
+        assert not torch.equal(a[n], c[n])
+
+
+def test_include_leaves_excluded_unchanged(device):
+    """Excluded params are bit-identical (not even shrunk); included change."""
+    model = torch.nn.Linear(8, 8).to(device)
+    w0 = model.weight.detach().clone()
+    b0 = model.bias.detach().clone()
+
+    shrink_and_perturb_(
+        model, shrink=0.5, perturb=0.1, include=lambda n, p: n == "weight"
+    )
+
+    assert torch.equal(model.bias, b0)
+    assert not torch.equal(model.weight, w0)
+
+
+def test_buffers_untouched(device):
+    """Buffers (batch-norm running stats) are never modified."""
+    model = torch.nn.BatchNorm1d(8).to(device)
+    with torch.no_grad():
+        model.running_mean.fill_(2.0)
+        model.running_var.fill_(3.0)
+    rm = model.running_mean.detach().clone()
+    rv = model.running_var.detach().clone()
+
+    shrink_and_perturb_(model, shrink=0.5, perturb=0.1)
+
+    assert torch.equal(model.running_mean, rm)
+    assert torch.equal(model.running_var, rv)
+    # weight (all ones -> std 0 -> no noise) is still shrunk.
+    assert torch.allclose(model.weight, torch.full_like(model.weight, 0.5))
+
+
+def test_invalid_arguments(device):
+    model = torch.nn.Linear(4, 4).to(device)
+    with pytest.raises(ValueError):
+        shrink_and_perturb_(model, shrink=-0.1)
+    with pytest.raises(ValueError):
+        shrink_and_perturb_(model, perturb=-1.0)
+    with pytest.raises(ValueError):
+        shrink_and_perturb_(model, noise="bogus")

From e739caed12efbb55533d05dbd4a38146c7a231e4 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Tue, 14 Jul 2026 09:26:46 -0600
Subject: [PATCH 02/10] fix: update noise interface and remove docs additions

---
 docs/api/nn/layers/weight_init.rst         |  6 ------
 docs/api/physicsnemo.nn.layers.rst         |  1 -
 physicsnemo/nn/module/utils/weight_init.py | 17 ++++++++++-------
 test/nn/module/test_weight_init.py         |  4 +++-
 4 files changed, 13 insertions(+), 15 deletions(-)
 delete mode 100644 docs/api/nn/layers/weight_init.rst

diff --git a/docs/api/nn/layers/weight_init.rst b/docs/api/nn/layers/weight_init.rst
deleted file mode 100644
index a387191198..0000000000
--- a/docs/api/nn/layers/weight_init.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-Weight Initialization
-=====================
-
-.. automodule:: physicsnemo.nn.module.utils.weight_init
-   :members:
-   :show-inheritance:
diff --git a/docs/api/physicsnemo.nn.layers.rst b/docs/api/physicsnemo.nn.layers.rst
index 04122df807..e472934730 100644
--- a/docs/api/physicsnemo.nn.layers.rst
+++ b/docs/api/physicsnemo.nn.layers.rst
@@ -18,4 +18,3 @@ PhysicsNeMo Layers
    nn/layers/regularization
    nn/layers/specialized
    nn/layers/graph_geometry
-   nn/layers/weight_init
diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 87910ddf98..4d4a316635 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -49,7 +49,7 @@ def shrink_and_perturb_(
     shrink: float = 0.5,
     perturb: float = 0.1,
     *,
-    noise: Union[str, Callable[[str, torch.Tensor], torch.Tensor]] = "scaled_normal",
+    noise: Union[str, Callable[[torch.Tensor], torch.Tensor]] = "scaled_normal",
     include: Optional[Callable[[str, torch.Tensor], bool]] = None,
     generator: Optional[torch.Generator] = None,
 ) -> torch.nn.Module:
@@ -94,9 +94,12 @@ def shrink_and_perturb_(
           elements per parameter tensor).
         - ``"normal"``: :math:`\varepsilon = z \sim \mathcal{N}(0, 1)`,
           unscaled.
-        - a callable ``(name, param) -> tensor`` returning
-          :math:`\varepsilon`, e.g. ``lambda n, p: fresh[n]`` to interpolate
-          toward the weights of a freshly constructed reference model.
+        - a callable ``(param) -> tensor`` returning :math:`\varepsilon`, with
+          the same signature as ``torch.randn_like``. Pass ``torch.randn_like``
+          (equivalent to ``"normal"``), the ``randn_like`` of a
+          ``StackedRandomGenerator``, any ``torch.nn.init``-style sampler, or
+          e.g. ``lambda p: torch.rand_like(p) * 2 - 1`` for a different noise
+          distribution.
     include : callable, optional
         Predicate ``(name, param) -> bool`` selecting which parameters to
         perturb. Parameters for which it returns ``False`` are left entirely
@@ -136,13 +139,13 @@ def shrink_and_perturb_(
         noise_fn = noise
     elif noise == "scaled_normal":
 
-        def noise_fn(name: str, p: torch.Tensor) -> torch.Tensor:
+        def noise_fn(p: torch.Tensor) -> torch.Tensor:
             z = torch.empty_like(p).normal_(generator=generator)
             return p.detach().std() * z
 
     elif noise == "normal":
 
-        def noise_fn(name: str, p: torch.Tensor) -> torch.Tensor:
+        def noise_fn(p: torch.Tensor) -> torch.Tensor:
             return torch.empty_like(p).normal_(generator=generator)
 
     else:
@@ -154,7 +157,7 @@ def noise_fn(name: str, p: torch.Tensor) -> torch.Tensor:
         if include is not None and not include(name, p):
             continue
         # eps is computed from the pre-shrink weight (matters for "scaled_normal").
-        eps = noise_fn(name, p)
+        eps = noise_fn(p)
         p.mul_(shrink).add_(eps, alpha=perturb)
     return module
 
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
index 9aad446e3b..5c5bb12b13 100644
--- a/test/nn/module/test_weight_init.py
+++ b/test/nn/module/test_weight_init.py
@@ -49,7 +49,9 @@ def test_shrink_and_perturb_callable_noise_exact(device):
     orig = {n: p.detach().clone() for n, p in model.named_parameters()}
     eps = {n: torch.ones_like(p) for n, p in model.named_parameters()}
 
-    shrink_and_perturb_(model, shrink=0.5, perturb=0.2, noise=lambda n, p: eps[n])
+    shrink_and_perturb_(
+        model, shrink=0.5, perturb=0.2, noise=lambda p: torch.ones_like(p)
+    )
 
     for n, p in model.named_parameters():
         assert torch.allclose(p, 0.5 * orig[n] + 0.2 * eps[n])

From b7259b1ec24c0400142976997484acd85944bcf9 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Tue, 14 Jul 2026 11:32:09 -0600
Subject: [PATCH 03/10] fix: guard scalar params from NaN and validate noise
 callable shape

---
 physicsnemo/nn/module/utils/weight_init.py | 30 ++++++++++++++++------
 test/nn/module/test_weight_init.py         | 29 +++++++++++++++++++++
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 4d4a316635..a4c9999774 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -90,16 +90,21 @@ def shrink_and_perturb_(
           :math:`\varepsilon = \operatorname{std}(\theta)\,z` with
           :math:`z \sim \mathcal{N}(0, 1)`, i.e. Gaussian noise scaled by the
           per-tensor standard deviation of the pre-shrink weight. Scale aware
-          and free of any architectural assumptions (expects at least two
-          elements per parameter tensor).
+          and free of any architectural assumptions. A scalar parameter (a
+          single element) has no defined spread and is shrunk only, without
+          added noise.
         - ``"normal"``: :math:`\varepsilon = z \sim \mathcal{N}(0, 1)`,
           unscaled.
         - a callable ``(param) -> tensor`` returning :math:`\varepsilon`, with
-          the same signature as ``torch.randn_like``. Pass ``torch.randn_like``
-          (equivalent to ``"normal"``), the ``randn_like`` of a
-          ``StackedRandomGenerator``, any ``torch.nn.init``-style sampler, or
-          e.g. ``lambda p: torch.rand_like(p) * 2 - 1`` for a different noise
+          the same signature as ``torch.randn_like``; it must return a tensor
+          matching the parameter's shape. Pass ``torch.randn_like`` (equivalent
+          to ``"normal"``), the ``randn_like`` of a ``StackedRandomGenerator``,
+          any ``torch.nn.init``-style sampler, or e.g.
+          ``lambda p: torch.rand_like(p) * 2 - 1`` for a different noise
           distribution.
+
+        The built-in presets sample Gaussian noise and therefore require
+        floating-point (or complex) parameters.
     include : callable, optional
         Predicate ``(name, param) -> bool`` selecting which parameters to
         perturb. Parameters for which it returns ``False`` are left entirely
@@ -119,8 +124,9 @@ def shrink_and_perturb_(
     Raises
     ------
     ValueError
-        If ``shrink`` or ``perturb`` is negative, or ``noise`` is an unknown
-        string.
+        If ``shrink`` or ``perturb`` is negative, ``noise`` is an unknown
+        string, or a ``noise`` callable returns a tensor whose shape does not
+        match its parameter.
 
     Examples
     --------
@@ -140,6 +146,9 @@ def shrink_and_perturb_(
     elif noise == "scaled_normal":
 
         def noise_fn(p: torch.Tensor) -> torch.Tensor:
+            if p.numel() < 2:
+                # std is undefined for a single element; shrink only, no noise.
+                return torch.zeros_like(p)
             z = torch.empty_like(p).normal_(generator=generator)
             return p.detach().std() * z
 
@@ -158,6 +167,11 @@ def noise_fn(p: torch.Tensor) -> torch.Tensor:
             continue
         # eps is computed from the pre-shrink weight (matters for "scaled_normal").
         eps = noise_fn(p)
+        if eps.shape != p.shape:
+            raise ValueError(
+                f"noise for parameter '{name}' returned shape "
+                f"{tuple(eps.shape)}, expected {tuple(p.shape)}"
+            )
         p.mul_(shrink).add_(eps, alpha=perturb)
     return module
 
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
index 5c5bb12b13..83b83e9b0b 100644
--- a/test/nn/module/test_weight_init.py
+++ b/test/nn/module/test_weight_init.py
@@ -20,6 +20,15 @@
 from physicsnemo.nn import shrink_and_perturb_
 
 
+class _ScalarParamNet(torch.nn.Module):
+    """Module with a scalar (numel==1) parameter, e.g. a learnable temperature."""
+
+    def __init__(self):
+        super().__init__()
+        self.fc = torch.nn.Linear(8, 8)
+        self.scale = torch.nn.Parameter(torch.tensor(0.07))
+
+
 def test_shrink_and_perturb_shrink_only(device):
     """perturb=0 shrinks every param exactly; returns the same module."""
     model = torch.nn.Linear(8, 8).to(device)
@@ -127,6 +136,26 @@ def test_buffers_untouched(device):
     assert torch.allclose(model.weight, torch.full_like(model.weight, 0.5))
 
 
+def test_scalar_param_shrink_only(device):
+    """A scalar (numel==1) param is shrunk without NaN under scaled_normal."""
+    model = _ScalarParamNet().to(device)
+    scale0 = model.scale.detach().clone()
+
+    shrink_and_perturb_(model, shrink=0.5, perturb=0.1)  # default scaled_normal
+
+    for p in model.parameters():
+        assert torch.isfinite(p).all()
+    # scalar has no defined spread -> shrink only, no noise.
+    assert torch.allclose(model.scale, 0.5 * scale0)
+
+
+def test_callable_wrong_shape_raises(device):
+    """A noise callable returning a mismatched shape is rejected, not broadcast."""
+    model = torch.nn.Linear(4, 4).to(device)
+    with pytest.raises(ValueError):
+        shrink_and_perturb_(model, noise=lambda p: torch.tensor(1.0, device=p.device))
+
+
 def test_invalid_arguments(device):
     model = torch.nn.Linear(4, 4).to(device)
     with pytest.raises(ValueError):

From f3332adb7a404354fad4a10277b534856b5a71f7 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Tue, 14 Jul 2026 12:49:20 -0600
Subject: [PATCH 04/10] fix: harden noise-callable contract, tied-param filter,
 and non-finite guards

---
 physicsnemo/nn/module/utils/weight_init.py | 118 +++++++++++++++++----
 test/nn/module/test_weight_init.py         |  93 ++++++++++++++++
 2 files changed, 189 insertions(+), 22 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index a4c9999774..61f3dc32d9 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -14,8 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 import warnings
-from typing import Callable, Optional, Union
+from typing import Callable, Literal, Optional, Union
 
 import numpy as np
 import torch
@@ -49,7 +50,9 @@ def shrink_and_perturb_(
     shrink: float = 0.5,
     perturb: float = 0.1,
     *,
-    noise: Union[str, Callable[[torch.Tensor], torch.Tensor]] = "scaled_normal",
+    noise: Union[
+        Literal["scaled_normal", "normal"], Callable[[torch.Tensor], torch.Tensor]
+    ] = "scaled_normal",
     include: Optional[Callable[[str, torch.Tensor], bool]] = None,
     generator: Optional[torch.Generator] = None,
 ) -> torch.nn.Module:
@@ -95,12 +98,13 @@ def shrink_and_perturb_(
           added noise.
         - ``"normal"``: :math:`\varepsilon = z \sim \mathcal{N}(0, 1)`,
           unscaled.
-        - a callable ``(param) -> tensor`` returning :math:`\varepsilon`, with
-          the same signature as ``torch.randn_like``; it must return a tensor
-          matching the parameter's shape. Pass ``torch.randn_like`` (equivalent
-          to ``"normal"``), the ``randn_like`` of a ``StackedRandomGenerator``,
-          any ``torch.nn.init``-style sampler, or e.g.
-          ``lambda p: torch.rand_like(p) * 2 - 1`` for a different noise
+        - a callable producing :math:`\varepsilon`. It is handed *fresh scratch
+          storage* carrying the parameter's shape, dtype, device, and layout
+          (never the live parameter, so it can neither alias nor overwrite the
+          weight) and must return a tensor of that shape. It may fill the given
+          tensor in place (e.g. ``torch.nn.init.normal_``) or return a new one
+          (e.g. ``torch.randn_like``, equivalent to ``"normal"``);
+          ``lambda p: torch.rand_like(p) * 2 - 1`` gives a different
           distribution.
 
         The built-in presets sample Gaussian noise and therefore require
@@ -110,7 +114,10 @@ def shrink_and_perturb_(
         perturb. Parameters for which it returns ``False`` are left entirely
         unchanged (not even shrunk). Default: all parameters. For warm-starting,
         pass e.g. ``include=lambda n, p: n in transferred`` to perturb only the
-        transferred backbone.
+        transferred backbone. Tied parameters are visited under every alias so
+        the predicate can match any checkpoint key, but each underlying tensor
+        is updated at most once. A warning is issued if an explicit ``include``
+        matches no parameters.
     generator : torch.Generator, optional
         Generator for the built-in Gaussian noise, for reproducibility. Must be
         on the same device as ``module``'s parameters. Ignored when ``noise``
@@ -124,37 +131,72 @@ def shrink_and_perturb_(
     Raises
     ------
     ValueError
-        If ``shrink`` or ``perturb`` is negative, ``noise`` is an unknown
-        string, or a ``noise`` callable returns a tensor whose shape does not
-        match its parameter.
+        If ``shrink`` or ``perturb`` is negative or non-finite, ``noise`` is an
+        unknown string, or a ``noise`` callable returns a tensor whose shape or
+        device is incompatible with its parameter (or complex noise for a real
+        parameter).
+
+    Notes
+    -----
+    Intended warm-start workflow: load the pretrained weights, apply this to the
+    transferred parameters, then (re)create the optimizer, and only afterwards
+    ``torch.compile`` / wrap in DDP or FSDP. Run it *before* distributed
+    wrapping: under FSDP2 the per-tensor ``std`` of ``"scaled_normal"`` reduces
+    over a sharded parameter and triggers an all-gather.
+
+    ``"scaled_normal"`` is a deliberately model-agnostic variant: it scales the
+    noise by the *pretrained tensor's* own standard deviation rather than by a
+    freshly re-initialized network's per-layer initializer variance (the form
+    in Ash & Adams, §4). It needs no architecture knowledge and matches the
+    validated production recipe.
+
+    The operation is not transactional across parameters: each parameter's noise
+    is validated before that parameter is modified, but if a later parameter
+    raises, earlier ones are already updated.
 
     Examples
     --------
     >>> import torch
     >>> from physicsnemo.nn import shrink_and_perturb_
     >>> model = torch.nn.Linear(4, 4)
+    >>> model.load_state_dict(pretrained_state)  # doctest: +SKIP
     >>> _ = shrink_and_perturb_(model, shrink=0.6, perturb=0.1)
     """
-    if shrink < 0.0 or perturb < 0.0:
+    if (
+        not math.isfinite(shrink)
+        or not math.isfinite(perturb)
+        or shrink < 0.0
+        or perturb < 0.0
+    ):
         raise ValueError(
-            f"shrink and perturb must be non-negative, got shrink={shrink}, "
-            f"perturb={perturb}"
+            f"shrink and perturb must be finite and non-negative, got "
+            f"shrink={shrink}, perturb={perturb}"
         )
 
+    # Resolve the noise source. Presets read the live (pre-shrink) parameter --
+    # only to *read* its std and allocate fresh noise, never to mutate it. A
+    # user callable instead receives fresh scratch storage carrying the
+    # parameter's shape/dtype/device/layout, so it can neither alias nor
+    # overwrite the weight (supporting both out-of-place samplers like
+    # ``torch.randn_like`` and in-place ones like ``torch.nn.init.normal_``).
     if callable(noise):
-        noise_fn = noise
+        user_noise = noise
+
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
+            return user_noise(torch.empty_like(p))
+
     elif noise == "scaled_normal":
 
-        def noise_fn(p: torch.Tensor) -> torch.Tensor:
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
             if p.numel() < 2:
                 # std is undefined for a single element; shrink only, no noise.
                 return torch.zeros_like(p)
             z = torch.empty_like(p).normal_(generator=generator)
-            return p.detach().std() * z
+            return z.mul_(p.detach().std())  # scale in place: one temporary
 
     elif noise == "normal":
 
-        def noise_fn(p: torch.Tensor) -> torch.Tensor:
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
             return torch.empty_like(p).normal_(generator=generator)
 
     else:
@@ -162,17 +204,49 @@ def noise_fn(p: torch.Tensor) -> torch.Tensor:
             f'Invalid noise "{noise}"; expected a callable or one of {_NOISE_PRESETS}'
         )
 
-    for name, p in module.named_parameters():
+    # Visit tied parameters under every alias so `include` can match any
+    # checkpoint key, but update each underlying tensor at most once.
+    seen: set[int] = set()
+    matched = False
+    for name, p in module.named_parameters(remove_duplicate=False):
         if include is not None and not include(name, p):
             continue
-        # eps is computed from the pre-shrink weight (matters for "scaled_normal").
-        eps = noise_fn(p)
+        matched = True
+        if id(p) in seen:
+            continue
+        seen.add(id(p))
+
+        if perturb == 0.0:
+            # Pure shrink (a no-op at shrink == 1); never draw from the generator.
+            if shrink != 1.0:
+                p.mul_(shrink)
+            continue
+
+        eps = make_eps(p)
+        # Validate before mutating so a rejected noise value leaves this
+        # parameter untouched (the op is not transactional across parameters).
         if eps.shape != p.shape:
             raise ValueError(
                 f"noise for parameter '{name}' returned shape "
                 f"{tuple(eps.shape)}, expected {tuple(p.shape)}"
             )
+        if eps.device != p.device:
+            raise ValueError(
+                f"noise for parameter '{name}' is on device {eps.device}, "
+                f"expected {p.device}"
+            )
+        if torch.is_complex(eps) and not torch.is_complex(p):
+            raise ValueError(
+                f"noise for parameter '{name}' is complex but the parameter is real"
+            )
         p.mul_(shrink).add_(eps, alpha=perturb)
+
+    if include is not None and not matched:
+        warnings.warn(
+            "shrink_and_perturb_: `include` matched no parameters; "
+            "nothing was modified.",
+            stacklevel=2,
+        )
     return module
 
 
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
index 83b83e9b0b..b92327e588 100644
--- a/test/nn/module/test_weight_init.py
+++ b/test/nn/module/test_weight_init.py
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 import pytest
 import torch
 
@@ -29,6 +31,16 @@ def __init__(self):
         self.scale = torch.nn.Parameter(torch.tensor(0.07))
 
 
+class _TiedNet(torch.nn.Module):
+    """Module with a tied weight shared under two names (embed / head)."""
+
+    def __init__(self):
+        super().__init__()
+        self.embed = torch.nn.Linear(8, 8, bias=False)
+        self.head = torch.nn.Linear(8, 8, bias=False)
+        self.head.weight = self.embed.weight
+
+
 def test_shrink_and_perturb_shrink_only(device):
     """perturb=0 shrinks every param exactly; returns the same module."""
     model = torch.nn.Linear(8, 8).to(device)
@@ -156,6 +168,77 @@ def test_callable_wrong_shape_raises(device):
         shrink_and_perturb_(model, noise=lambda p: torch.tensor(1.0, device=p.device))
 
 
+def test_inplace_init_callable(device):
+    """In-place torch.nn.init samplers fill scratch, never alias the weight."""
+    model = torch.nn.Linear(4, 4).to(device)
+    orig = {n: p.detach().clone() for n, p in model.named_parameters()}
+
+    shrink_and_perturb_(model, shrink=0.5, perturb=0.2, noise=torch.nn.init.ones_)
+
+    for n, p in model.named_parameters():
+        # eps == ones; if the callable had aliased the live param the pretrained
+        # value would be lost (would collapse to a constant).
+        assert torch.allclose(p, 0.5 * orig[n] + 0.2 * torch.ones_like(p))
+
+
+def test_complex_noise_into_real_raises(device):
+    """Complex noise for a real parameter is rejected before any mutation."""
+    model = torch.nn.Linear(4, 4).to(device)
+    w0 = model.weight.detach().clone()
+    with pytest.raises(ValueError):
+        shrink_and_perturb_(
+            model, noise=lambda p: torch.ones_like(p, dtype=torch.complex64)
+        )
+    assert torch.equal(model.weight, w0)  # untouched: validated before mutate
+
+
+def test_tied_params_selected_by_alias(device):
+    """A tied weight is reachable via any alias name and updated exactly once."""
+    model = _TiedNet().to(device)
+    w0 = model.embed.weight.detach().clone()
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")  # no "matched no parameters" warning
+        shrink_and_perturb_(
+            model, shrink=0.5, perturb=0.0, include=lambda n, p: n == "head.weight"
+        )
+
+    # Updated once via the alias -> 0.5*w0 (0.25*w0 would mean double application).
+    assert torch.allclose(model.embed.weight, 0.5 * w0)
+
+
+def test_include_no_match_warns(device):
+    model = torch.nn.Linear(4, 4).to(device)
+    with pytest.warns(UserWarning, match="matched no parameters"):
+        shrink_and_perturb_(model, include=lambda n, p: False)
+
+
+def test_normal_magnitude(device):
+    """'normal' noise has unit std, so with shrink=0 the result std ~ perturb."""
+    model = torch.nn.Linear(256, 256).to(device)
+    gen = torch.Generator(device=device).manual_seed(0)
+    shrink_and_perturb_(
+        model,
+        shrink=0.0,
+        perturb=0.1,
+        noise="normal",
+        include=lambda n, p: n == "weight",
+        generator=gen,
+    )
+    assert abs(model.weight.std().item() - 0.1) / 0.1 < 0.1
+
+
+def test_zero_perturb_preserves_rng(device):
+    """perturb=0 must not advance the global RNG (no noise is drawn)."""
+    model = torch.nn.Linear(8, 8).to(device)
+    torch.manual_seed(123)
+    a = torch.randn(4, device=device)
+    torch.manual_seed(123)
+    shrink_and_perturb_(model, shrink=0.5, perturb=0.0)
+    b = torch.randn(4, device=device)
+    assert torch.equal(a, b)
+
+
 def test_invalid_arguments(device):
     model = torch.nn.Linear(4, 4).to(device)
     with pytest.raises(ValueError):
@@ -164,3 +247,13 @@ def test_invalid_arguments(device):
         shrink_and_perturb_(model, perturb=-1.0)
     with pytest.raises(ValueError):
         shrink_and_perturb_(model, noise="bogus")
+
+
+def test_nonfinite_coefficients_raise(device):
+    """NaN / +-inf coefficients are rejected before touching the module."""
+    model = torch.nn.Linear(4, 4).to(device)
+    for bad in (float("nan"), float("inf"), float("-inf")):
+        with pytest.raises(ValueError):
+            shrink_and_perturb_(model, shrink=bad)
+        with pytest.raises(ValueError):
+            shrink_and_perturb_(model, perturb=bad)

From dc81f26dd1e942c9fa72e885cc990176e3184346 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Tue, 14 Jul 2026 13:40:51 -0600
Subject: [PATCH 05/10] docs: update changelog

---
 CHANGELOG.md | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4f04773beb..cbc6011ddc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,11 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Adds `physicsnemo.nn.shrink_and_perturb_`, an in-place shrink-and-perturb
-  weight re-initialization (Ash & Adams, NeurIPS 2020) for warm-starting from
-  pretrained weights. Model-agnostic, with a per-tensor scale-aware Gaussian
-  default (`theta <- shrink*theta + perturb*std(theta)*noise`), a plain-Gaussian
-  and custom-callable noise source, and an `include` filter to restrict the
-  perturbation to (e.g.) transferred backbone parameters.
+  weight re-initialization for warm-starting from pretrained weights.
 - Adds a `global_shape` argument to `ShardTensor.from_local`, enabling the
   no-communication `sharding_shapes="chunk"` path.
 - Adds exact-boundary quality mesh generation to

From 3b2cedee350afd56cf756f1183d27222ce52c805 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Tue, 14 Jul 2026 14:22:41 -0600
Subject: [PATCH 06/10] refactor: extract noise resolution and validation
 helpers in shrink_and_perturb_

---
 physicsnemo/nn/module/utils/weight_init.py | 110 ++++++++++++---------
 test/nn/module/test_weight_init.py         |  33 +++----
 2 files changed, 77 insertions(+), 66 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 61f3dc32d9..3f1055cac3 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -44,6 +44,70 @@ def trunc_normal_(*args, **kwargs):
 _NOISE_PRESETS = ("scaled_normal", "normal")
 
 
+def _resolve_noise(
+    noise: Union[
+        Literal["scaled_normal", "normal"], Callable[[torch.Tensor], torch.Tensor]
+    ],
+    generator: Optional[torch.Generator],
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Build the ``(param) -> noise`` sampler for :func:`shrink_and_perturb_`.
+
+    Presets read the live (pre-shrink) parameter only to *read* its std and
+    allocate fresh noise, never to mutate it. A user callable instead receives
+    fresh scratch storage carrying the parameter's shape/dtype/device/layout, so
+    it can neither alias nor overwrite the weight (supporting both out-of-place
+    samplers like ``torch.randn_like`` and in-place ones like
+    ``torch.nn.init.normal_``).
+    """
+    if callable(noise):
+
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
+            return noise(torch.empty_like(p))
+
+    elif noise == "scaled_normal":
+
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
+            if p.numel() < 2:
+                # std is undefined for a single element; shrink only, no noise.
+                return torch.zeros_like(p)
+            z = torch.empty_like(p).normal_(generator=generator)
+            return z.mul_(p.detach().std())  # scale in place: one temporary
+
+    elif noise == "normal":
+
+        def make_eps(p: torch.Tensor) -> torch.Tensor:
+            return torch.empty_like(p).normal_(generator=generator)
+
+    else:
+        raise ValueError(
+            f'Invalid noise "{noise}"; expected a callable or one of {_NOISE_PRESETS}'
+        )
+
+    return make_eps
+
+
+def _validate_noise(eps: torch.Tensor, p: torch.Tensor, name: str) -> None:
+    """Reject noise incompatible with parameter ``name`` before it is applied.
+
+    Checking up front keeps each parameter's update atomic: a bad value raises
+    without having partially modified the parameter.
+    """
+    if eps.shape != p.shape:
+        raise ValueError(
+            f"noise for parameter '{name}' returned shape "
+            f"{tuple(eps.shape)}, expected {tuple(p.shape)}"
+        )
+    if eps.device != p.device:
+        raise ValueError(
+            f"noise for parameter '{name}' is on device {eps.device}, "
+            f"expected {p.device}"
+        )
+    if torch.is_complex(eps) and not torch.is_complex(p):
+        raise ValueError(
+            f"noise for parameter '{name}' is complex but the parameter is real"
+        )
+
+
 @torch.no_grad()
 def shrink_and_perturb_(
     module: torch.nn.Module,
@@ -173,36 +237,7 @@ def shrink_and_perturb_(
             f"shrink={shrink}, perturb={perturb}"
         )
 
-    # Resolve the noise source. Presets read the live (pre-shrink) parameter --
-    # only to *read* its std and allocate fresh noise, never to mutate it. A
-    # user callable instead receives fresh scratch storage carrying the
-    # parameter's shape/dtype/device/layout, so it can neither alias nor
-    # overwrite the weight (supporting both out-of-place samplers like
-    # ``torch.randn_like`` and in-place ones like ``torch.nn.init.normal_``).
-    if callable(noise):
-        user_noise = noise
-
-        def make_eps(p: torch.Tensor) -> torch.Tensor:
-            return user_noise(torch.empty_like(p))
-
-    elif noise == "scaled_normal":
-
-        def make_eps(p: torch.Tensor) -> torch.Tensor:
-            if p.numel() < 2:
-                # std is undefined for a single element; shrink only, no noise.
-                return torch.zeros_like(p)
-            z = torch.empty_like(p).normal_(generator=generator)
-            return z.mul_(p.detach().std())  # scale in place: one temporary
-
-    elif noise == "normal":
-
-        def make_eps(p: torch.Tensor) -> torch.Tensor:
-            return torch.empty_like(p).normal_(generator=generator)
-
-    else:
-        raise ValueError(
-            f'Invalid noise "{noise}"; expected a callable or one of {_NOISE_PRESETS}'
-        )
+    make_eps = _resolve_noise(noise, generator)
 
     # Visit tied parameters under every alias so `include` can match any
     # checkpoint key, but update each underlying tensor at most once.
@@ -225,20 +260,7 @@ def make_eps(p: torch.Tensor) -> torch.Tensor:
         eps = make_eps(p)
         # Validate before mutating so a rejected noise value leaves this
         # parameter untouched (the op is not transactional across parameters).
-        if eps.shape != p.shape:
-            raise ValueError(
-                f"noise for parameter '{name}' returned shape "
-                f"{tuple(eps.shape)}, expected {tuple(p.shape)}"
-            )
-        if eps.device != p.device:
-            raise ValueError(
-                f"noise for parameter '{name}' is on device {eps.device}, "
-                f"expected {p.device}"
-            )
-        if torch.is_complex(eps) and not torch.is_complex(p):
-            raise ValueError(
-                f"noise for parameter '{name}' is complex but the parameter is real"
-            )
+        _validate_noise(eps, p, name)
         p.mul_(shrink).add_(eps, alpha=perturb)
 
     if include is not None and not matched:
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
index b92327e588..f977468575 100644
--- a/test/nn/module/test_weight_init.py
+++ b/test/nn/module/test_weight_init.py
@@ -78,25 +78,29 @@ def test_shrink_and_perturb_callable_noise_exact(device):
         assert torch.allclose(p, 0.5 * orig[n] + 0.2 * eps[n])
 
 
-def test_scaled_normal_magnitude(device):
-    """scaled_normal noise std tracks perturb * std(theta) per tensor."""
+@pytest.mark.parametrize(
+    "noise, w_init_std",  # scaled_normal scales noise by std(theta); normal does not
+    [("scaled_normal", 3.0), ("normal", 1.0)],
+)
+def test_noise_magnitude(device, noise, w_init_std):
+    """With shrink=0 the result std ~ perturb * (std(theta) for scaled_normal, else 1)."""
     model = torch.nn.Linear(256, 256).to(device)
     with torch.no_grad():
-        model.weight.normal_(0.0, 3.0)
-    w_std = model.weight.std().item()
+        model.weight.normal_(0.0, w_init_std)
+    scale = model.weight.std().item() if noise == "scaled_normal" else 1.0
 
     gen = torch.Generator(device=device).manual_seed(0)
-    # shrink=0 isolates the noise term: result == perturb * std(theta) * z.
+    # shrink=0 isolates the noise term: result == perturb * scale * z.
     shrink_and_perturb_(
         model,
         shrink=0.0,
         perturb=0.1,
-        noise="scaled_normal",
+        noise=noise,
         include=lambda n, p: n == "weight",
         generator=gen,
     )
 
-    expected = 0.1 * w_std
+    expected = 0.1 * scale
     assert abs(model.weight.std().item() - expected) / expected < 0.1
 
 
@@ -213,21 +217,6 @@ def test_include_no_match_warns(device):
         shrink_and_perturb_(model, include=lambda n, p: False)
 
 
-def test_normal_magnitude(device):
-    """'normal' noise has unit std, so with shrink=0 the result std ~ perturb."""
-    model = torch.nn.Linear(256, 256).to(device)
-    gen = torch.Generator(device=device).manual_seed(0)
-    shrink_and_perturb_(
-        model,
-        shrink=0.0,
-        perturb=0.1,
-        noise="normal",
-        include=lambda n, p: n == "weight",
-        generator=gen,
-    )
-    assert abs(model.weight.std().item() - 0.1) / 0.1 < 0.1
-
-
 def test_zero_perturb_preserves_rng(device):
     """perturb=0 must not advance the global RNG (no noise is drawn)."""
     model = torch.nn.Linear(8, 8).to(device)

From a8e28d63832b190795fe2d96c497e0dc273bdc65 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Mon, 20 Jul 2026 12:11:06 -0600
Subject: [PATCH 07/10] feat: complete reference for s&p

---
 physicsnemo/nn/module/utils/weight_init.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 3f1055cac3..444a978b60 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -133,6 +133,7 @@ def shrink_and_perturb_(
     the pretrained features. In warm-started training this often reaches a lower
     loss asymptote than fine-tuning the raw pretrained weights (Ash & Adams,
     *On Warm-Starting Neural Network Training*, NeurIPS 2020).
+    https://arxiv.org/pdf/1910.08475
 
     The operation is only meaningful on **pretrained** weights: applied to a
     fresh initialization it merely rescales and re-noises random values.
@@ -210,8 +211,8 @@ def shrink_and_perturb_(
 
     ``"scaled_normal"`` is a deliberately model-agnostic variant: it scales the
     noise by the *pretrained tensor's* own standard deviation rather than by a
-    freshly re-initialized network's per-layer initializer variance (the form
-    in Ash & Adams, §4). It needs no architecture knowledge and matches the
+    freshly re-initialized network's per-layer initializer variance.
+    It needs no architecture knowledge and matches the
     validated production recipe.
 
     The operation is not transactional across parameters: each parameter's noise

From 7ea00186d48ab5e8d989970433d408a482d161a2 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Mon, 20 Jul 2026 12:20:53 -0600
Subject: [PATCH 08/10] fix: parameter init to randn_like for readability

---
 physicsnemo/nn/module/utils/weight_init.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 444a978b60..258c7af366 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -70,13 +70,13 @@ def make_eps(p: torch.Tensor) -> torch.Tensor:
             if p.numel() < 2:
                 # std is undefined for a single element; shrink only, no noise.
                 return torch.zeros_like(p)
-            z = torch.empty_like(p).normal_(generator=generator)
+            z = torch.randn_like(p, generator=generator)
             return z.mul_(p.detach().std())  # scale in place: one temporary
 
     elif noise == "normal":
 
         def make_eps(p: torch.Tensor) -> torch.Tensor:
-            return torch.empty_like(p).normal_(generator=generator)
+            return torch.randn_like(p, generator=generator)
 
     else:
         raise ValueError(

From 7bd8b0907f1c63d03c17684ef21837165708707a Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Mon, 20 Jul 2026 12:24:51 -0600
Subject: [PATCH 09/10] fix: check generator device

---
 physicsnemo/nn/module/utils/weight_init.py | 19 ++++++++++++++++---
 test/nn/module/test_weight_init.py         | 13 +++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 258c7af366..48ccab6ef6 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -197,9 +197,10 @@ def shrink_and_perturb_(
     ------
     ValueError
         If ``shrink`` or ``perturb`` is negative or non-finite, ``noise`` is an
-        unknown string, or a ``noise`` callable returns a tensor whose shape or
-        device is incompatible with its parameter (or complex noise for a real
-        parameter).
+        unknown string, a preset ``generator`` is on a different device than the
+        module's parameters, or a ``noise`` callable returns a tensor whose
+        shape or device is incompatible with its parameter (or complex noise for
+        a real parameter).
 
     Notes
     -----
@@ -238,6 +239,18 @@ def shrink_and_perturb_(
             f"shrink={shrink}, perturb={perturb}"
         )
 
+    # The preset samplers draw on `generator`, which must sit on the device of
+    # every parameter it seeds (a callable ignores it, and perturb == 0 draws
+    # nothing, so skip the check in those cases).
+    if generator is not None and not callable(noise) and perturb != 0.0:
+        bad = {p.device for p in module.parameters()} - {generator.device}
+        if bad:
+            raise ValueError(
+                f"generator is on {generator.device}, but the module has "
+                f"parameters on {sorted(map(str, bad))}; the generator must "
+                f"match the parameter device(s)"
+            )
+
     make_eps = _resolve_noise(noise, generator)
 
     # Visit tied parameters under every alias so `include` can match any
diff --git a/test/nn/module/test_weight_init.py b/test/nn/module/test_weight_init.py
index f977468575..d03c19fba0 100644
--- a/test/nn/module/test_weight_init.py
+++ b/test/nn/module/test_weight_init.py
@@ -246,3 +246,16 @@ def test_nonfinite_coefficients_raise(device):
             shrink_and_perturb_(model, shrink=bad)
         with pytest.raises(ValueError):
             shrink_and_perturb_(model, perturb=bad)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="device mismatch needs a second device"
+)
+def test_generator_device_mismatch_raises():
+    """A preset generator on the wrong device is rejected up front."""
+    model = torch.nn.Linear(4, 4).cuda()
+    cpu_gen = torch.Generator()  # cpu, mismatched with the cuda parameters
+    with pytest.raises(ValueError, match="generator is on"):
+        shrink_and_perturb_(model, generator=cpu_gen)
+    # a callable ignores the generator, so a device mismatch must not raise
+    shrink_and_perturb_(model, generator=cpu_gen, noise=torch.randn_like)

From 0611a41752cf5b451f542332274720f441e18a46 Mon Sep 17 00:00:00 2001
From: Carmelo Gonzales <43048528+melo-gonzo@users.noreply.github.com>
Date: Mon, 20 Jul 2026 13:29:20 -0600
Subject: [PATCH 10/10] fix: raise NotImplementedError for models that are
 sharded

---
 physicsnemo/nn/module/utils/weight_init.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/physicsnemo/nn/module/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
index 48ccab6ef6..143f7f7ab8 100644
--- a/physicsnemo/nn/module/utils/weight_init.py
+++ b/physicsnemo/nn/module/utils/weight_init.py
@@ -20,6 +20,7 @@
 
 import numpy as np
 import torch
+from torch.distributed.tensor import DTensor
 from torch.nn.init import trunc_normal_ as _torch_trunc_normal_
 
 
@@ -142,7 +143,10 @@ def shrink_and_perturb_(
     ----------
     module : torch.nn.Module
         Module whose parameters are perturbed in place. Buffers (e.g.
-        batch-norm running statistics) are left untouched.
+        batch-norm running statistics) are left untouched. Only whole, unsharded
+        parameters are supported: a sharded ``DTensor`` parameter (e.g. from
+        FSDP2 ``fully_shard`` or tensor parallelism) raises
+        ``NotImplementedError``, so apply this before distributed wrapping.
     shrink : float, optional
         Multiplicative retention factor :math:`\lambda \ge 0` applied to each
         weight. Values in ``[0, 1)`` shrink toward zero; ``1.0`` disables the
@@ -201,14 +205,16 @@ def shrink_and_perturb_(
         module's parameters, or a ``noise`` callable returns a tensor whose
         shape or device is incompatible with its parameter (or complex noise for
         a real parameter).
+    NotImplementedError
+        If any selected parameter is a sharded ``DTensor`` (only whole,
+        unsharded tensors are supported).
 
     Notes
     -----
     Intended warm-start workflow: load the pretrained weights, apply this to the
     transferred parameters, then (re)create the optimizer, and only afterwards
     ``torch.compile`` / wrap in DDP or FSDP. Run it *before* distributed
-    wrapping: under FSDP2 the per-tensor ``std`` of ``"scaled_normal"`` reduces
-    over a sharded parameter and triggers an all-gather.
+    wrapping.
 
     ``"scaled_normal"`` is a deliberately model-agnostic variant: it scales the
     noise by the *pretrained tensor's* own standard deviation rather than by a
@@ -239,6 +245,13 @@ def shrink_and_perturb_(
             f"shrink={shrink}, perturb={perturb}"
         )
 
+    if any(isinstance(p, DTensor) for p in module.parameters()):
+        raise NotImplementedError(
+            "shrink_and_perturb_ does not support sharded DTensor parameters "
+            "(e.g. from FSDP2 `fully_shard` or tensor parallelism); apply it to "
+            "the full model before distributed wrapping."
+        )
+
     # The preset samplers draw on `generator`, which must sit on the device of
     # every parameter it seeds (a callable ignores it, and perturb == 0 draws
     # nothing, so skip the check in those cases).