Merge branch 'main' into psharpe/GLOBE-perf-to-merge

Author: peterdsharpe

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -31,7 +31,7 @@ body:
     attributes:
       label: Version
       description: What version of PhysicsNeMo are you running?
-      placeholder: "example: 2.0.0"
+      placeholder: "example: 2.1.0"
     validations:
       required: true
 

CHANGELOG.md

@@ -6,7 +6,7 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [2.1.0a0] - 2026-XX-YY
+## [2.2.0] - 2026-XX-YY
 
 ### Added
 
@@ -21,15 +21,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (`SpatialBranch`), and coordinate features.
 - Adds `Sin` elementwise sine activation to `physicsnemo.nn`, registered
   in `ACT2FN` so it can be looked up by name (`get_activation("sin")`).
-- Adds GLOBE model (`physicsnemo.experimental.models.globe.model.GLOBE`),
-  including new variant that uses a dual tree traversal algorithm to reduce the
-  complexity of the kernel evaluations from O(N^2) to O(N).
-- Adds GLOBE AirFRANS example case (`examples/cfd/external_aerodynamics/globe/airfrans`)
-- Adds GLOBE DrivAerML example case (`examples/cfd/external_aerodynamics/globe/drivaer`)
-- Adds drop-test dynamics recipe.
-- Adds concrete dropout uncertainty quantification for GeoTransolver. Learnable
-  per-layer dropout rates enable MC-Dropout inference for uncertainty
-  estimates. Disabled by default (`concrete_dropout: false`).
 - Adds active-learning recipe for external-aerodynamics surrogates
   (`examples/cfd/external_aerodynamics/active_learning_aero/`). Iteratively
   fine-tunes a GP-augmented GeoTransolver onto an out-of-distribution
@@ -39,6 +30,57 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   protocols and `physicsnemo.experimental.uq.VariationalGPHead`, with a
   layered structure (generic AL driver / GP-UQ recipe / aero adapter)
   designed for reuse on other UQ-based regression problems.
+
+### Changed
+
+### Deprecated
+
+### Removed
+
+### Fixed
+
+- Replaced three plain-string regex / docstring literals containing invalid
+  escape sequences with raw-string equivalents
+  (`physicsnemo/utils/logging/launch.py`,
+  `physicsnemo/metrics/general/calibration.py`,
+  `physicsnemo/metrics/general/crps.py`); these were `SyntaxWarning`s today
+  and become `SyntaxError`s in Python 3.16.
+- Various test cleanups to remove self-inflicted warnings in CI output:
+  disabled pytest collection for `TestModelA`/`TestModelB` helpers in
+  `test/core/test_registry.py` via `__test__ = False`; migrated
+  `test/nn/module/test_interpolation.py` to call the non-deprecated
+  `grid_to_point_interpolation` and added a dedicated test for the
+  deprecation alias; scoped a `lr_scheduler.step()`-before-`optimizer.step()`
+  `UserWarning` filter to a single test in
+  `test/optim/test_combined_optimizer.py`; guarded the
+  `DistributedManager.initialize()` calls in `test/utils/test_checkpoint.py`
+  with `is_initialized()`; and suppressed the import-time
+  `ExperimentalFeatureWarning` in `test/datapipes/healda/test_features.py`
+  via `warnings.catch_warnings()`.
+- Fixed `physicsnemo.utils.get_checkpoint_dir` returning paths with `\`
+  separators on Windows (e.g. `.\checkpoints_model`), which was inconsistent
+  with the `/`-based paths used elsewhere in the checkpoint utilities and
+  broke the `test_get_checkpoint_dir` CI test on Windows. The function now
+  always joins with `/`, working uniformly for local paths and `fsspec`
+  URIs (`msc://`, etc.) across operating systems.
+
+### Security
+
+### Dependencies
+
+## [2.1.0] - 2026-05-26
+
+### Added
+
+- Adds GLOBE model (`physicsnemo.experimental.models.globe.model.GLOBE`),
+  including new variant that uses a dual tree traversal algorithm to reduce the
+  complexity of the kernel evaluations from O(N^2) to O(N).
+- Adds GLOBE AirFRANS example case (`examples/cfd/external_aerodynamics/globe/airfrans`)
+- Adds GLOBE DrivAerML example case (`examples/cfd/external_aerodynamics/globe/drivaer`)
+- Adds drop-test dynamics recipe.
+- Adds concrete dropout uncertainty quantification for GeoTransolver. Learnable
+  per-layer dropout rates enable MC-Dropout inference for uncertainty
+  estimates. Disabled by default (`concrete_dropout: false`).
 - Adds automatic support for `FSDP` and/or `ShardTensor` models in checkpoint save/load
   functionality
 - PhysicsNeMo-Mesh now supports conversion from PyVista/VTK/VTU meshes that may
@@ -226,42 +268,46 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   combined-workflow and from-checkpoint round-trip tests. Most tests
   run with `fullgraph=True` and `error_on_recompile` to catch
   `torch.compile` regressions.
+- Internal weight initialization in the distributed AFNO layers and the
+  `EarthAttention` blocks of `physicsnemo.nn.module.attention_layers` now
+  dispatches to `torch.nn.init.trunc_normal_` directly instead of going
+  through frozen in-tree copies of the pre-PyTorch-2.12 inverse-CDF
+  implementation. PyTorch 2.12 reimplemented `trunc_normal_` as a
+  rejection-sampling loop on top of `normal_()` (see
+  [pytorch/pytorch#174997](https://github.com/pytorch/pytorch/pull/174997)),
+  so seeded from-scratch initialization consumes the RNG stream
+  differently on 2.12+ vs older versions. Existing trained checkpoints
+  are unaffected (loading bypasses init). Forward-accuracy reference
+  outputs for `AFNO`, `ModAFNO`, `Transolver`, `FLARE`, and `Pangu` were
+  regenerated against the new algorithm. Rather than wiring per-model
+  skips, `test.common.validate_forward_accuracy` now uniformly skips on
+  `torch < 2.12` (the reference data is locked to that floor via a single
+  `_REFERENCE_DATA_MIN_TORCH` constant; bump it when a PyTorch
+  release next changes an init/RNG algorithm any forward-accuracy model
+  depends on, and regenerate the `.pth` files at the same time).
 
 ### Deprecated
 
 - `physicsnemo.utils.mesh` is deprecated and will be removed in v2.2.0. For
   isosurface extraction, use `physicsnemo.mesh.generate.marching_cubes` instead
   of `sdf_to_stl`. For VTP/OBJ/STL file conversion (`combine_vtp_files`,
   `convert_tesselated_files_in_directory`), use VTK or PyVista directly.
+- `physicsnemo.nn.module.utils.trunc_normal_` (and its submodule path
+  `physicsnemo.nn.module.utils.weight_init.trunc_normal_`) is deprecated
+  and will be removed in v2.2.0. It is now a thin wrapper around
+  `torch.nn.init.trunc_normal_` that emits a `DeprecationWarning` on
+  call, replacing the frozen in-tree copy of the legacy inverse-CDF
+  implementation. Use `torch.nn.init.trunc_normal_` directly.
 
 ### Removed
 
+- The legacy in-tree `trunc_normal_` implementation that lived in
+  `physicsnemo/models/afno/distributed/layers.py` (`_trunc_normal_` /
+  `_no_grad_trunc_normal_`) is removed. These names were private; all
+  in-tree call sites now use `torch.nn.init.trunc_normal_`.
+
 ### Fixed
 
-- Replaced three plain-string regex / docstring literals containing invalid
-  escape sequences with raw-string equivalents
-  (`physicsnemo/utils/logging/launch.py`,
-  `physicsnemo/metrics/general/calibration.py`,
-  `physicsnemo/metrics/general/crps.py`); these were `SyntaxWarning`s today
-  and become `SyntaxError`s in Python 3.16.
-- Various test cleanups to remove self-inflicted warnings in CI output:
-  disabled pytest collection for `TestModelA`/`TestModelB` helpers in
-  `test/core/test_registry.py` via `__test__ = False`; migrated
-  `test/nn/module/test_interpolation.py` to call the non-deprecated
-  `grid_to_point_interpolation` and added a dedicated test for the
-  deprecation alias; scoped a `lr_scheduler.step()`-before-`optimizer.step()`
-  `UserWarning` filter to a single test in
-  `test/optim/test_combined_optimizer.py`; guarded the
-  `DistributedManager.initialize()` calls in `test/utils/test_checkpoint.py`
-  with `is_initialized()`; and suppressed the import-time
-  `ExperimentalFeatureWarning` in `test/datapipes/healda/test_features.py`
-  via `warnings.catch_warnings()`.
-- Fixed `physicsnemo.utils.get_checkpoint_dir` returning paths with `\`
-  separators on Windows (e.g. `.\checkpoints_model`), which was inconsistent
-  with the `/`-based paths used elsewhere in the checkpoint utilities and
-  broke the `test_get_checkpoint_dir` CI test on Windows. The function now
-  always joins with `/`, working uniformly for local paths and `fsspec`
-  URIs (`msc://`, etc.) across operating systems.
 - Fixed functional benchmark plot fallback labeling so unlabeled ASV results use
   the same key ordering as the benchmark runner.
 - Fixed graph break caused by `FunctionSpec` dispatch (`max(key=)` is not supported by `torch.compile`)
@@ -305,12 +351,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed the sinusoidal positional embeddings formula in `SongUNet` and
   `MultiDiffusionModel2D` so it now follows the standard `sin / cos`
   convention. Affected reference data was regenerated.
-
-### Security
+- Constructing a `Mesh` (or `DomainMesh`) inside a `torch.compile`-traced
+  function no longer raises `AttributeError` / `KeyError` or silently
+  produces wrong output. The breakage came from two regressions in
+  `tensordict >= 0.12.0` (PR `pytorch/tensordict#1552`), where the
+  `@tensorclass` init wrapper's bypass branch silently skipped both
+  field-default normalization and `__post_init__` under
+  `torch.compile`. We pin `tensordict < 0.12` until the upstream fix
+  (`pytorch/tensordict#1708`, `pytorch/tensordict#1709`) ships, and add
+  a regression test (`test/mesh/mesh/test_compile.py`) that constructs
+  a `Mesh` inside `torch.compile` and reads cached properties, so the
+  same bug cannot return on a future pin bump unnoticed.
 
 ### Dependencies
 
 - Increments minimum viable PyTorch version to `torch>=2.5.0` to support FSDP better
+- Upper-bounds `tensordict < 0.12` to avoid the `torch.compile` regressions
+  in `tensordict >= 0.12.0` (see corresponding entry under Fixed).
 
 ## [2.0.0] - 2026-03-09
 

Dockerfile

@@ -26,7 +26,9 @@ FROM ${BASE_CONTAINER} AS builder
 ARG TARGETPLATFORM
 
 # Install uv (use system Python for installs; set so --system is default)
-COPY --from=ghcr.io/astral-sh/uv:0.10.3 /uv /uvx /bin/
+# Pinned to 0.11.14 (latest stable as of May 2026) which bundles
+# rustls-webpki >= 0.103.13 (fixes GHSA-82j2-j2ch-gfr8).
+COPY --from=ghcr.io/astral-sh/uv:0.11.14 /uv /uvx /bin/
 ENV UV_SYSTEM_PYTHON=1
 # Base image Python is PEP 668 externally-managed; allow system installs in container
 ENV UV_BREAK_SYSTEM_PACKAGES=1

examples/additive_manufacturing/sintering_physics/README.md

@@ -58,13 +58,25 @@ For more sample parts simulation:
 
 - Download PhysicsNeMo, make or install
 
+- Install the example's Python dependencies:
+
+  ```bash
+  pip install -r requirements.txt
+  ```
+
 - Find the matching torch-scatter version with torch and cuda version enabled:
   - i.e. pip install torch-scatter-f `https://data.pyg.org/whl/torch-2.2.0%2Bcu121/torch_scatter-2.1.2%2Bpt22cu121-cp311-cp311-linux_x86_64.whl`
     (replace the torch-scatter wheel with the matching cuda, torch version )
   - torch-scatter installation guide: `https://pypi.org/project/torch-scatter/`
   - wheels source: `https://data.pyg.org/whl/`
 
-- pip install tensorflow
+- Install TensorFlow separately (used by `reading_utils.py`,
+  `graph_dataset.py`, and the `data_process/` scripts for
+  `tf.train.SequenceExample` I/O and `tf.data` pipelines):
+
+  ```bash
+  pip install "tensorflow>=2.15,<3.0"
+  ```
 
   - test version: tensorflow-2.15.0.post1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 

examples/additive_manufacturing/sintering_physics/inference.py

@@ -32,15 +32,6 @@
         "Mesh Graph Net Datapipe requires the Tensorflow library. Install the "
         + "package at: https://www.tensorflow.org/install"
     )
-physical_devices = tf.config.list_physical_devices("GPU")
-
-try:
-    for device_ in physical_devices:
-        tf.config.experimental.set_memory_growth(device_, True)
-except:
-    # Invalid device or cannot modify virtual devices once initialized.
-    pass
-
 import hydra
 import torch
 from graph_dataset import GraphDataset

examples/additive_manufacturing/sintering_physics/requirements.txt

@@ -1,3 +1,12 @@
-# pyvista is optional, required if need to run data preprocessing from raw simulation
-# pyvista==0.32.1
-tensorflow>=2.15,<3.0 # generate tfrecord
+# Core dependencies for training, inference, and rollout rendering
+dm-tree
+tqdm
+hydra-core
+omegaconf
+matplotlib
+pyvista
+vtk
+natsort
+scikit-learn
+tensorboard
+tensorflow-cpu>=2.15,<3.0

examples/additive_manufacturing/sintering_physics/train.py

@@ -58,14 +58,6 @@
 )
 from physicsnemo.models.vfgn.graph_network_modules import VFGNLearnedSimulator
 
-physical_devices = tf.config.list_physical_devices("GPU")
-try:
-    for device_ in physical_devices:
-        tf.config.experimental.set_memory_growth(device_, True)
-except:
-    # Invalid device or cannot modify virtual devices once initialized.
-    pass
-
 
 def Train(rank_zero_logger, dist, cfg: DictConfig):
     """
@@ -131,6 +123,7 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
     writer = SummaryWriter(log_dir=cfg.data_options.ckpt_path_vfgn)
 
     optimizer = None
+    scaler = None
     # todo : check device
     device = "cpu"
     step = 0
@@ -178,18 +171,23 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
 
         sampled_noise *= noise_mask
 
-        pred_target = model(
-            next_positions=targets.to(device),
-            position_sequence=inputs.to(device),
-            position_sequence_noise=sampled_noise.to(device),
-            n_particles_per_example=features["n_particles_per_example"].to(device),
-            n_edges_per_example=features["n_edges_per_example"].to(device),
-            senders=features["senders"].to(device),
-            receivers=features["receivers"].to(device),
-            predict_length=cfg.train_options.pred_len,
-            particle_types=features["particle_type"].to(device),
-            global_context=features.get("step_context").to(device),
-        )
+        amp_enabled = cfg.general.fp16 and scaler is not None
+        with torch.autocast(
+            device_type=device.type if isinstance(device, torch.device) else "cpu",
+            enabled=amp_enabled,
+        ):
+            pred_target = model(
+                next_positions=targets.to(device),
+                position_sequence=inputs.to(device),
+                position_sequence_noise=sampled_noise.to(device),
+                n_particles_per_example=features["n_particles_per_example"].to(device),
+                n_edges_per_example=features["n_edges_per_example"].to(device),
+                senders=features["senders"].to(device),
+                receivers=features["receivers"].to(device),
+                predict_length=cfg.train_options.pred_len,
+                particle_types=features["particle_type"].to(device),
+                global_context=features.get("step_context").to(device),
+            )
 
         if optimizer is None:
             # first data need to inference the feature size
@@ -208,14 +206,7 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
             model = model.to(device)
             optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
             if cfg.general.fp16:
-                # double check if amp installed
-                try:
-                    from apex import amp
-
-                    model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
-                except ImportError as e:
-                    print("Apex package not available -> ", e)
-                    exit()
+                scaler = torch.amp.GradScaler(device.type)
 
             scheduler = torch.optim.lr_scheduler.ExponentialLR(
                 optimizer, gamma=0.1, verbose=True
@@ -394,11 +385,12 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
         # back propogation
         optimizer.zero_grad()
         if cfg.general.fp16:
-            with amp.scale_loss(loss, optimizer) as scaled_loss:
-                scaled_loss.backward()
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
         else:
             loss.backward()
-        optimizer.step()
+            optimizer.step()
 
         running_loss += loss.item()
 

examples/cfd/darcy_fno/requirements.txt

@@ -1,2 +1,3 @@
 hydra-core>=1.2.0
 termcolor>=2.1.1
+matplotlib

examples/cfd/darcy_nested_fnos/requirements.txt

@@ -1 +1 @@
-mlflow>=2.1.1
+mlflow>=3.12.0

examples/cfd/darcy_physics_informed/requirements.txt

@@ -1 +1,3 @@
-gdown
+gdown>=5.2.2
+matplotlib
+scipy