Add optional extra diffusion times and outputs

.importlinter

@@ -117,7 +117,8 @@ layers =
    samplers | metrics
    multi_diffusion | preconditioners | guidance
    noise_schedulers
-   utils | base
+   utils
+   base
 
 [importlinter:contract:physicsnemo-sym]
 name = Control Internal Dependencies in PhysicsNeMo sym

CHANGELOG.md

@@ -15,6 +15,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Adds an experimental, independently implemented SMART operator with sampled
   geometry encoding, reusable latent contexts, and chunked field decoding.
 
+- Adds `DiffusionExtraTimes` and `DiffusionExtraOutputs` to
+  `physicsnemo.diffusion` for utilities that carry named times or auxiliary
+  outputs. Compatible model interfaces and preconditioners can accept
+  `(t, extra_times)`, where `extra_times` is a `DiffusionExtraTimes` container,
+  and return `(prediction, extra_outputs)`, where `extra_outputs` is a
+  `DiffusionExtraOutputs` container. The reserved `r` field supports a
+  destination time as in MeanFlow. `logvar` supports learned loss weighting
+  as in sCM and EDM2. Existing single-time, single-prediction calls remain
+  unchanged. Existing DSM and flow-matching losses continue to consume
+  single tensors.
 - Adds standalone FLARE++ attention and model APIs, with input-conditioned
   dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside
   GeoTransolver.

docs/api/diffusion/introduction.rst

@@ -290,6 +290,16 @@ API Reference
     :members:
     :exclude-members: __init__
 
+:code:`DiffusionExtraTimes`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: physicsnemo.diffusion.DiffusionExtraTimes
+
+:code:`DiffusionExtraOutputs`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: physicsnemo.diffusion.DiffusionExtraOutputs
+
 :code:`Predictor`
 ~~~~~~~~~~~~~~~~~
 

docs/api/diffusion/models.rst

@@ -80,6 +80,63 @@ expects ``(x, t, condition)``.  A short adapter class solves this:
         def forward(self, x, t, condition=None):
             return self.net(x, noise_labels=t, class_labels=condition)
 
+Most diffusion and flow-matching models take one time tensor and return one
+prediction tensor.
+A compatible :class:`~physicsnemo.diffusion.DiffusionModel` can instead
+accept ``(t, extra_times)`` with
+:class:`~physicsnemo.diffusion.DiffusionExtraTimes` and return
+``(prediction, extra_outputs)`` with
+:class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. For example, MeanFlow
+uses a destination time ``r``, while sCM and EDM2 can use ``logvar`` for loss
+weighting. The primary time and prediction remain required tensors.
+
+The adapter below combines the destination time with the primary time for a
+DiT backbone. A small
+:class:`~physicsnemo.models.mlp.FullyConnected` head predicts a log-scale
+from the primary time. This example does not define a canonical MeanFlow
+architecture. A compatible loss decides how to use ``logvar`` during
+training. Inference code that needs only a prediction can ignore it.
+
+.. code-block:: python
+
+    import torch
+    from physicsnemo.core import Module
+    from physicsnemo.diffusion import DiffusionExtraOutputs, DiffusionExtraTimes
+    from physicsnemo.models.dit import DiT
+    from physicsnemo.models.mlp import FullyConnected
+
+    class ExtendedDiTAdapter(Module):
+        def __init__(self):
+            super().__init__()
+            self.backbone = DiT(
+                input_size=8, in_channels=2, out_channels=2,
+                patch_size=4, hidden_size=32, depth=1, num_heads=4,
+            )
+            self.r_embed = FullyConnected(
+                in_features=1, layer_size=8, out_features=1, num_layers=1,
+            )
+            self.logvar_head = FullyConnected(
+                in_features=1, layer_size=8, out_features=1, num_layers=1,
+            )
+
+        def forward(self, x, t, condition=None):
+            t, extra_times = t
+            # DiT takes one time tensor, so combine the two times in the adapter.
+            r_embed = self.r_embed(extra_times.r[:, None]).squeeze(-1)
+            prediction = self.backbone(x, t + r_embed, condition)
+            logvar = self.logvar_head(t[:, None]).squeeze(-1)
+            return prediction, DiffusionExtraOutputs(
+                logvar=logvar, batch_size=[x.shape[0]],
+            )
+
+    model = ExtendedDiTAdapter()
+    x = torch.randn(2, 2, 8, 8)
+    t = torch.ones(2)
+    times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2]))
+    prediction, extra_outputs = model(x, times)
+    assert prediction.shape == x.shape
+    assert extra_outputs.logvar.shape == t.shape
+
 
 
 .. _backbone_external_models:
@@ -100,4 +157,3 @@ You can bring in models from external libraries, such as:
 * by using :meth:`~physicsnemo.core.module.Module.from_torch` to convert a
 ``torch.nn.Module`` subclass into a :class:`~physicsnemo.core.module.Module`
 with PhysicsNeMo serialization support
-

physicsnemo/diffusion/__init__.py

@@ -14,4 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base import Denoiser, DiffusionModel, Predictor, PredictorType  # noqa: F401
+from .base import (  # noqa: F401
+    Denoiser,
+    DiffusionExtraOutputs,
+    DiffusionExtraTimes,
+    DiffusionModel,
+    Predictor,
+    PredictorType,
+)