diff --git a/.importlinter b/.importlinter
index fa39832fae..891ae686a7 100644
--- a/.importlinter
+++ b/.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
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8cb26a3be9..b76c9207ab 100644
--- a/CHANGELOG.md
+++ b/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.
diff --git a/docs/api/diffusion/introduction.rst b/docs/api/diffusion/introduction.rst
index 15863eb88f..06993b70f4 100644
--- a/docs/api/diffusion/introduction.rst
+++ b/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`
 ~~~~~~~~~~~~~~~~~
 
diff --git a/docs/api/diffusion/models.rst b/docs/api/diffusion/models.rst
index 1ae6ac54d0..927a9f9b32 100644
--- a/docs/api/diffusion/models.rst
+++ b/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
-
diff --git a/physicsnemo/diffusion/__init__.py b/physicsnemo/diffusion/__init__.py
index d9602b252f..7b7d2bf629 100644
--- a/physicsnemo/diffusion/__init__.py
+++ b/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,
+)
diff --git a/physicsnemo/diffusion/base.py b/physicsnemo/diffusion/base.py
index f386ba3a32..f2acfe093b 100644
--- a/physicsnemo/diffusion/base.py
+++ b/physicsnemo/diffusion/base.py
@@ -19,11 +19,166 @@
 from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
-from jaxtyping import Float
-from tensordict import TensorDict
+from jaxtyping import Float, Shaped
+from tensordict import TensorClass, TensorDict
+
+
+class DiffusionExtraTimes(TensorClass):
+    r"""Named, batched diffusion times accompanying the primary time tensor.
+
+    This ``tensordict.TensorClass`` container carries named extra times for
+    diffusion utilities. A :class:`DiffusionModel` is one possible consumer.
+
+    Most diffusion and flow-matching models take a single diffusion time
+    tensor. Models that need more than one diffusion time during training
+    or sampling can take ``(t, extra_times)`` instead. The required primary
+    time ``t`` stays separate from the optional named fields in this
+    container.
+
+    Every supplied time tensor has shape :math:`(B,)` and uses the same
+    scheduler coordinates as the primary time. Set ``batch_size=[B]`` to
+    match the input batch. For example, utilities reserve ``r`` for the
+    destination diffusion time in MeanFlow. Each field is independently
+    optional and defaults to ``None``.
+
+    .. note::
+
+        Subclass this container to add time fields for a custom use. Do not
+        reuse reserved names for different meanings. The framework does not
+        infer the meaning of custom fields; the consuming utility must define
+        and handle them.
+
+    Parameters
+    ----------
+    r : torch.Tensor or None, optional, default=None
+        Destination diffusion time for MeanFlow, with shape :math:`(B,)`.
+        Leave unset when the model does not use this field.
+    batch_size : sequence of int
+        Set to ``[B]``, matching the primary time tensor.
+
+    Examples
+    --------
+    Pair a primary time with the reserved destination time for a MeanFlow
+    consumer. The batch sizes of both tensors and the container agree.
+    This example shows the tuple passed to a compatible utility:
+
+    >>> import torch
+    >>> from physicsnemo.diffusion import DiffusionExtraTimes
+    >>> t = torch.ones(2)
+    >>> # The reserved r field holds the destination time.
+    >>> times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2]))
+    >>> t, extra_times = times
+    >>> t.shape, extra_times.r.shape
+    (torch.Size([2]), torch.Size([2]))
+
+    A subclass can add independent optional times for a custom consumer.
+    The consumer must interpret these fields and can accept either field
+    alone or both together:
+
+    >>> from jaxtyping import Float
+    >>> class CustomTimes(DiffusionExtraTimes):
+    ...     # The consuming utility defines these custom times.
+    ...     r1: Float[torch.Tensor, " B"] | None = None
+    ...     r2: Float[torch.Tensor, " B"] | None = None
+    >>> first = CustomTimes(r1=torch.zeros(2), batch_size=[2])
+    >>> second = CustomTimes(r2=torch.ones(2), batch_size=[2])
+    >>> first.r is None, first.r2 is None, second.r1 is None
+    (True, True, True)
+    >>> first.r1.shape, second.r2.shape
+    (torch.Size([2]), torch.Size([2]))
+    >>> both = CustomTimes(
+    ...     r1=torch.zeros(2), r2=torch.ones(2), batch_size=[2]
+    ... )
+    >>> both.r1.shape, both.r2.shape
+    (torch.Size([2]), torch.Size([2]))
+    """
+
+    r: Float[torch.Tensor, " B"] | None = None
+
+
+class DiffusionExtraOutputs(TensorClass):
+    r"""Named, batched auxiliary outputs accompanying a model prediction.
+
+    This ``tensordict.TensorClass`` container carries named extra outputs for
+    diffusion utilities. A :class:`DiffusionModel` is one possible producer.
+
+    Most diffusion and flow-matching models return a single tensor that
+    predicts clean data, noise, a score, or a velocity, as described by
+    :data:`PredictorType`. Models that produce auxiliary outputs for
+    training or sampling can return ``(prediction, extra_outputs)`` instead.
+    For example, sCM and EDM2 losses use a learned log-scale alongside the
+    prediction for adaptive loss weighting. The required prediction stays
+    separate from the optional named fields in this container.
+
+    All supplied output tensors share the prediction's batch size
+    :math:`B`, but their remaining dimensions can differ. Set
+    ``batch_size=[B]`` to describe that shared batch dimension. For example,
+    the reserved field ``logvar`` provides a learned log-scale for adaptive
+    loss weighting in sCM and EDM2. It does not represent diffusion
+    transition variance. Each field is independently optional and defaults
+    to ``None``.
+
+    .. note::
+
+        Subclass this container to add outputs for a custom use. Do not reuse
+        reserved names for different meanings. The framework does not infer
+        the meaning of custom fields; the consuming utility must define and
+        handle them.
+
+    Parameters
+    ----------
+    logvar : torch.Tensor or None, optional, default=None
+        Learned log-scale for adaptive loss weighting in sCM and EDM2, with
+        shape :math:`(B,)`. Leave unset when the model does not use it.
+    batch_size : sequence of int
+        Set to ``[B]``, matching the prediction's batch dimension.
+
+    Examples
+    --------
+    Pair a required prediction with a learned log-scale for a compatible
+    loss. The prediction retains its own shape while the container shares
+    its batch dimension:
+
+    >>> import torch
+    >>> from physicsnemo.diffusion import DiffusionExtraOutputs
+    >>> prediction = torch.zeros(2, 3, 5)
+    >>> # The reserved logvar field supplies loss weighting metadata.
+    >>> output = (
+    ...     prediction,
+    ...     DiffusionExtraOutputs(logvar=torch.zeros(2), batch_size=[2]),
+    ... )
+    >>> prediction, extra_outputs = output
+    >>> prediction.shape, extra_outputs.logvar.shape
+    (torch.Size([2, 3, 5]), torch.Size([2]))
+
+    A subclass can hold independent optional outputs with different feature
+    shapes. The consumer decides how to use each field, so these names carry
+    no framework-defined meaning:
+
+    >>> from jaxtyping import Float
+    >>> class CustomOutputs(DiffusionExtraOutputs):
+    ...     # The consuming utility defines these custom outputs.
+    ...     features: Float[torch.Tensor, " B C H W"] | None = None
+    ...     confidence: Float[torch.Tensor, " B K"] | None = None
+    >>> extra = CustomOutputs(
+    ...     features=torch.zeros(2, 4, 5, 7),
+    ...     confidence=torch.ones(2, 3),
+    ...     batch_size=[2],
+    ... )
+    >>> extra.features.shape, extra.confidence.shape
+    (torch.Size([2, 4, 5, 7]), torch.Size([2, 3]))
+    >>> extra.logvar is None
+    True
+    >>> features_only = CustomOutputs(features=extra.features, batch_size=[2])
+    >>> features_only.confidence is None
+    True
+    """
+
+    logvar: Float[torch.Tensor, " B"] | None = None
+
 
 PredictorType = Literal["x0", "score", "epsilon", "flow"]
-"""Type of prediction produced by a diffusion model.
+"""Type of prediction produced by a diffusion or flow-matching model.
 
 One of:
 
@@ -46,8 +201,18 @@ class DiffusionModel(Protocol):
     into a prediction. This protocol defines the standard interface that all
     such models must meet.
 
-    Any model or function that implements this interface can be used with
-    preconditioners, losses, samplers, and other diffusion utilities.
+    Most models take a single diffusion time tensor and return a single
+    prediction tensor. Specialized models can extend either side of this
+    interface:
+
+    - Pass ``(t, extra_times)`` with :class:`DiffusionExtraTimes` for models
+      that accept more than one diffusion time.
+    - Return ``(prediction, extra_outputs)`` with
+      :class:`DiffusionExtraOutputs` for models that produce auxiliary outputs.
+
+    The primary time and prediction remain required tensors, separate from
+    the optional named fields. Use tuples with consumers that support the
+    supplied fields. Prediction targets follow :data:`PredictorType`.
 
     The interface is **prediction-agnostic**: whether your model predicts
     clean data (:math:`\mathbf{x}_0`), noise (:math:`\epsilon`), score
@@ -69,6 +234,9 @@ class DiffusionModel(Protocol):
 
     Examples
     --------
+    An ordinary model can take one time tensor and return one prediction.
+    The protocol checks that the model exposes the expected call interface:
+
     >>> import torch
     >>> import torch.nn.functional as F
     >>> from physicsnemo.diffusion import DiffusionModel
@@ -79,15 +247,40 @@ class DiffusionModel(Protocol):
     ...
     >>> isinstance(Model(), DiffusionModel)
     True
+
+    A compatible model can also consume an extra time and return an auxiliary
+    output. This example keeps the prediction as the first tuple element and
+    uses the reserved fields with their intended meanings:
+
+    >>> from physicsnemo.diffusion import DiffusionExtraTimes, DiffusionExtraOutputs
+    >>> class ExtendedModel:
+    ...     def __call__(self, x, t, condition=None, **kwargs):
+    ...         t, extra_times = t
+    ...         prediction = x * (t - extra_times.r)[:, None]
+    ...         extra_outputs = DiffusionExtraOutputs(
+    ...             logvar=torch.zeros_like(t), batch_size=[x.shape[0]]
+    ...         )
+    ...         # The required prediction stays first in the output tuple.
+    ...         return prediction, extra_outputs
+    >>> x = torch.ones(2, 3)
+    >>> t = torch.ones(2)
+    >>> times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2]))
+    >>> prediction, extra_outputs = ExtendedModel()(x, times)
+    >>> prediction.shape, extra_outputs.logvar.shape
+    (torch.Size([2, 3]), torch.Size([2]))
     """
 
     def __call__(
         self,
         x: Float[torch.Tensor, " B *dims"],
-        t: Float[torch.Tensor, " B"],
+        t: Float[torch.Tensor, " B"]
+        | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]],
         condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None,
         **model_kwargs: Any,
-    ) -> Float[torch.Tensor, " B *dims"]:
+    ) -> (
+        Float[torch.Tensor, " B *dims"]
+        | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]]
+    ):
         r"""
         Forward pass of the diffusion model.
 
@@ -97,8 +290,10 @@ def __call__(
             Noisy latent state of shape :math:`(B, *)` where :math:`B` is the
             batch size and :math:`*` denotes any number of additional
             dimensions (e.g., channels and spatial dimensions).
-        t : torch.Tensor
-            Diffusion time or noise level tensor of shape :math:`(B,)`.
+        t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes]
+            Diffusion time for each sample, with shape :math:`(B,)`.
+            Pass ``(t, extra_times)`` to supply extra named times in
+            :class:`DiffusionExtraTimes`, each with shape :math:`(B,)`.
         condition : torch.Tensor, TensorDict, or None, optional, default=None
             Conditioning information for the model. If a Tensor or a TensorDict
             is passed, it should have batch size :math:`B` matching that of
@@ -109,8 +304,10 @@ def __call__(
 
         Returns
         -------
-        torch.Tensor
-            Model output with the same shape as ``x``.
+        torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs]
+            Prediction with the shape of ``x``, optionally paired with a
+            :class:`DiffusionExtraOutputs`. Every call returns a prediction
+            tensor; auxiliary tensors share its batch dimension.
         """
         ...
 
@@ -130,7 +327,7 @@ class Predictor(Protocol):
     prediction. A predictor can be a trained neural network, a guidance
     function (e.g., classifier-free guidance, DPS-style guidance), or any
     combination thereof.  The exact meaning of the output depends on the
-    predictor type and how it is used. Any callable that implements this
+    :data:`PredictorType` and its use. Any callable that implements this
     interface can be used as a predictor in sampling utilities.
 
     This protocol is typically used during inference. For training, which
@@ -163,8 +360,9 @@ class Predictor(Protocol):
 
     Examples
     --------
-    **Example 1:** Convert a trained conditional model into a predictor using
-    ``functools.partial``:
+    **Example 1:** A conditional model needs ``y`` to form its clean-data
+    estimate. Bind that condition with ``functools.partial`` so sampling code
+    can call the resulting predictor with only ``x`` and ``t``:
 
     >>> import torch
     >>> from functools import partial
@@ -184,9 +382,10 @@ class Predictor(Protocol):
     >>> isinstance(x0_predictor, Predictor)
     True
 
-    **Example 2:** Convert the x0-predictor above into a score-predictor
-    (using a simple EDM-like schedule where :math:`\sigma(t) = t` and
-    :math:`\alpha(t) = 1`):
+    **Example 2:** A scheduler may need a score rather than the clean-data
+    estimate produced above. For the illustrative schedule
+    :math:`\sigma(t) = t` and :math:`\alpha(t) = 1`, convert that estimate
+    with the score identity while keeping the same predictor call signature:
 
     >>> def x0_to_score(x0, x_t, t):
     ...     sigma_sq = t.view(-1, 1) ** 2
@@ -198,6 +397,24 @@ class Predictor(Protocol):
     >>>
     >>> isinstance(score_predictor, Predictor)
     True
+
+    **Example 3:** A model can return :class:`DiffusionExtraOutputs` for
+    training, such as a ``logvar`` field. Wrap the model to select its
+    prediction tensor for sampling. Standard denoiser factories require a
+    predictor that accepts one diffusion time tensor:
+
+    >>> from physicsnemo.diffusion import DiffusionExtraOutputs
+    >>> def model_with_extras(x, t, condition=None):
+    ...     return x / 2, DiffusionExtraOutputs(
+    ...         logvar=torch.zeros(x.shape[0]), batch_size=[x.shape[0]]
+    ...     )
+    >>> def prediction_only(x, t):
+    ...     prediction, _ = model_with_extras(x, t)
+    ...     return prediction
+    >>> prediction_only(torch.ones(2, 3), torch.ones(2)).shape
+    torch.Size([2, 3])
+    >>> isinstance(prediction_only, Predictor)
+    True
     """
 
     def __call__(
@@ -270,8 +487,10 @@ class Denoiser(Protocol):
 
     Examples
     --------
-    Manually creating a denoiser from an x0-predictor using a simple EDM-like
-    schedule (:math:`\sigma(t)=t`, :math:`\alpha(t)=1`):
+    This example builds a denoiser from a clean-data predictor for the
+    illustrative schedule :math:`\sigma(t)=t` and :math:`\alpha(t)=1`.
+    The returned tensor is the update term consumed by a solver. It uses one
+    time tensor, as required by the standard predictor and denoiser interfaces:
 
     >>> import torch
     >>> from physicsnemo.diffusion import Denoiser
diff --git a/physicsnemo/diffusion/preconditioners/preconditioners.py b/physicsnemo/diffusion/preconditioners/preconditioners.py
index 404240330b..eeac45dd74 100644
--- a/physicsnemo/diffusion/preconditioners/preconditioners.py
+++ b/physicsnemo/diffusion/preconditioners/preconditioners.py
@@ -17,16 +17,21 @@
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+from functools import partial
 from typing import Any, Tuple
 
 import torch
-from jaxtyping import Float
+from jaxtyping import Float, Shaped
 from tensordict import TensorDict
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.placement_types import Replicate
 
 from physicsnemo.core.meta import ModelMetaData
 from physicsnemo.core.module import Module
+from physicsnemo.diffusion.base import (
+    DiffusionExtraOutputs,
+    DiffusionExtraTimes,
+)
 
 
 def _ensure_plain_tensor(t: torch.Tensor) -> torch.Tensor:
@@ -65,8 +70,8 @@ class BaseAffinePreconditioner(Module, ABC):
     transformation.
 
     This class provides a standardized interface for implementing
-    preconditioners that use affine transformations of the model
-    input and output.
+    preconditioners that use affine transformations of the model input and
+    output.
 
     The preconditioner wraps a neural network model :math:`F` and applies
     a preconditioning formula to transform the network output to produce
@@ -90,7 +95,7 @@ class BaseAffinePreconditioner(Module, ABC):
     The wrapped model :math:`F` must be an instance of
     :class:`~physicsnemo.core.Module` that satisfies the
     :class:`~physicsnemo.diffusion.DiffusionModel` interface, with the
-    following signature:
+    following signature for a single diffusion time and prediction tensor:
 
     .. code-block:: python
 
@@ -101,6 +106,16 @@ class BaseAffinePreconditioner(Module, ABC):
             **model_kwargs: Any,
         ) -> torch.Tensor  # Shape: (B, *)
 
+    .. note::
+
+        Models can also accept ``(t, extra_times)`` with
+        :class:`~physicsnemo.diffusion.DiffusionExtraTimes` and return
+        ``(prediction, extra_outputs)`` with
+        :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`.
+        Every supplied time tensor uses the same time-to-noise mapping and
+        noise encoding as the primary time. The affine formula transforms
+        only the prediction tensor. Auxiliary outputs pass through unchanged.
+
     The preconditioner is agnostic to the prediction target of the wrapped
     model :math:`F`. The same preconditioning formula is applied regardless of
     whether the model is an :math:`\mathbf{x}_0`-predictor, an
@@ -128,8 +143,11 @@ class BaseAffinePreconditioner(Module, ABC):
     x : torch.Tensor
         Noisy latent state of shape :math:`(B, *)` where :math:`B` is the
         batch size and :math:`*` denotes any number of additional dimensions.
-    t : torch.Tensor
-        Diffusion time tensor of shape :math:`(B,)`.
+    t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes]
+        Diffusion time for each sample, with shape :math:`(B,)`. Pass
+        ``(t, extra_times)`` to supply extra named diffusion times in
+        :class:`~physicsnemo.diffusion.DiffusionExtraTimes`.
+        The preconditioner encodes every supplied time for the backbone.
     condition : torch.Tensor, TensorDict, or None, optional, default=None
         Single Tensor or a TensorDict containing conditioning tensors with
         batch size :math:`B` matching that of ``x``. Pass ``None`` for an
@@ -140,9 +158,10 @@ class BaseAffinePreconditioner(Module, ABC):
 
     Outputs
     -------
-    torch.Tensor
-        Preconditioned model output with the same shape as the original model
-        output.
+    torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs]
+        Preconditioned prediction with the shape of ``x``, optionally paired
+        with auxiliary outputs. The affine formula transforms only the
+        prediction; auxiliary outputs keep their values and gradients.
 
     Note
     ----
@@ -182,7 +201,8 @@ class BaseAffinePreconditioner(Module, ABC):
     method since :math:`\sigma(t) = t` (noise level and diffusion time are the
     same).
 
-    We first define a simple model to wrap:
+    First, define a small backbone with the required call signature. Its
+    convolution provides a prediction with the same shape as its input:
 
     >>> import torch
     >>> from tensordict import TensorDict
@@ -196,7 +216,8 @@ class BaseAffinePreconditioner(Module, ABC):
     ...     def forward(self, x, t, condition=None):
     ...         return self.net(x)
 
-    Now we define the EDM preconditioner:
+    The subclass below supplies the four EDM coefficients. It wraps the
+    backbone and keeps the original prediction shape:
 
     >>> from physicsnemo.diffusion.preconditioners import (
     ...     BaseAffinePreconditioner,
@@ -225,9 +246,9 @@ class BaseAffinePreconditioner(Module, ABC):
     >>> out.shape
     torch.Size([2, 3, 16, 16])
 
-    The following example shows how to override the :meth:`sigma` method to
-    implement a Variance Exploding (VE) preconditioner where
-    :math:`\sigma(t) = \sqrt{t}`.
+    A different schedule can override :meth:`sigma` while retaining the
+    preconditioner interface. This VE example uses
+    :math:`\sigma(t) = \sqrt{t}`:
 
     >>> class VEPreconditioner(BaseAffinePreconditioner):
     ...     def __init__(self, model):
@@ -256,17 +277,16 @@ class BaseAffinePreconditioner(Module, ABC):
 
     Some models in PhysicsNeMo have signatures that differ from the
     :class:`~physicsnemo.diffusion.DiffusionModel` interface. Below are
-    examples showing how to write thin wrappers to make them compatible
-    with preconditioners, including image-based conditioning via channel
+    examples showing thin wrappers that make them compatible with
+    preconditioners. They also show image conditioning through channel
     concatenation.
 
     **Example: Wrapping SongUNet**
 
     The :class:`~physicsnemo.models.diffusion_unets.SongUNet` model has
     the signature ``forward(x, noise_labels, class_labels, augment_labels)``.
-    We wrap it to match ``forward(x, t, condition)``, where ``condition``
-    contains both class labels (1D vector) and an image to concatenate
-    channel-wise:
+    We wrap it to match ``forward(x, t, condition)``. The condition carries
+    class labels as a vector and an image to concatenate along channels:
 
     >>> from physicsnemo.models.diffusion_unets import SongUNet
     >>> from physicsnemo.diffusion import DiffusionModel
@@ -309,8 +329,8 @@ class BaseAffinePreconditioner(Module, ABC):
 
     The :class:`~physicsnemo.models.dit.DiT` model has
     the signature ``forward(x, t, condition, ...)``. We wrap it to support
-    both image conditioning (via channel concatenation) and vector
-    conditioning:
+    both image conditioning through channel concatenation and vector
+    conditioning. The adapter passes the vector to DiT's ``condition`` input:
 
     >>> from physicsnemo.models.dit import DiT
     >>> class DiTWrapper(Module):
@@ -355,7 +375,8 @@ class BaseAffinePreconditioner(Module, ABC):
     is also passed as a separate argument) is common across several diffusion
     use-cases.
     Thus for convenience, we provide the wrapper class :class:`~physicsnemo.diffusion.utils.ConcatConditionWrapper`
-    to save you the trouble of writing your own wrapper for this common pattern:
+    to use this common pattern. The example pairs it with an EDM
+    preconditioner and checks the output shape:
 
     >>> from physicsnemo.diffusion.preconditioners import EDMPreconditioner
     >>> from physicsnemo.diffusion.utils import ConcatConditionWrapper
@@ -451,13 +472,34 @@ def _validate_input(self, x: torch.Tensor) -> None:
         tracing.  The default implementation is a no-op.
         """
 
+    def _encode_time(
+        self,
+        t: Float[torch.Tensor, " B"],
+        x: Float[torch.Tensor, " B *dims"],
+        c_noise: Float[torch.Tensor, " B *ones"] | None = None,
+    ) -> Float[torch.Tensor, " B"]:
+        r"""Encode a diffusion time, reusing its noise coefficient when available."""
+        if c_noise is None:
+            shape = (-1,) + (1,) * (x.ndim - 1)
+            sigma = self.sigma(t).reshape(shape)  # (B, 1, ...)
+            _, c_noise, _, _ = self.compute_coefficients(_ensure_plain_tensor(sigma))
+        c_noise = _ensure_plain_tensor(c_noise)
+        if isinstance(x, DTensor):
+            c_noise = _replicate_on_mesh(c_noise, x.device_mesh)
+        return c_noise.flatten()  # (B,)
+
     def forward(
         self,
-        x: torch.Tensor,
-        t: torch.Tensor,
-        condition: torch.Tensor | TensorDict | None = None,
+        x: Float[torch.Tensor, " B *dims"],
+        t: Float[torch.Tensor, " B"]
+        | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]],
+        condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None,
         **model_kwargs: Any,
-    ) -> torch.Tensor:
+    ) -> (
+        Float[torch.Tensor, " B *dims"]
+        | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]]
+    ):
+        t, extra_times = t if isinstance(t, tuple) else (t, None)
         if not torch.compiler.is_compiling():
             B = x.shape[0]
             if t.shape != (B,):
@@ -465,6 +507,11 @@ def forward(
                     f"Expected t to have shape ({B},) matching batch size of "
                     f"x, but got {t.shape}."
                 )
+            if extra_times is not None:
+                if extra_times.batch_size != torch.Size([B]):
+                    raise ValueError(
+                        f"Expected time batch size [{B}], got {extra_times.batch_size}."
+                    )
             if isinstance(condition, TensorDict):
                 if condition.batch_size and condition.batch_size[0] != B:
                     raise ValueError(
@@ -481,7 +528,7 @@ def forward(
 
         # Map time step to noise level via sigma method
         expected_shape = (-1,) + (1,) * (x.ndim - 1)
-        sigma_t = self.sigma(t).reshape(expected_shape)
+        sigma_t = self.sigma(t).reshape(expected_shape)  # (B, 1, ...)
 
         # Compute preconditioning coefficients
         # Unwrap to plain tensors so that element-wise arithmetic between
@@ -495,9 +542,7 @@ def forward(
         # tensors, then re-promote to replicated DTensors on x's mesh when
         # x is a ShardTensor so that element-wise arithmetic between
         # coefficients and data is type-compatible.
-        c_in, c_noise, c_out, c_skip = (
-            _ensure_plain_tensor(c) for c in (c_in, c_noise, c_out, c_skip)
-        )
+        c_in, c_out, c_skip = (_ensure_plain_tensor(c) for c in (c_in, c_out, c_skip))
 
         # Only manually promote coefficients for genuine (non-Shard) DTensor
         # activations. A ShardTensor subclasses torch.Tensor (not DTensor), so
@@ -507,28 +552,32 @@ def forward(
         # otherwise breaks ops like PositionalEmbedding's torch.outer).
         if isinstance(x, DTensor):
             x_mesh = x.device_mesh
-            c_in, c_noise, c_out, c_skip = (
-                _replicate_on_mesh(c, x_mesh) for c in (c_in, c_noise, c_out, c_skip)
+            c_in, c_out, c_skip = (
+                _replicate_on_mesh(c, x_mesh) for c in (c_in, c_out, c_skip)
             )
 
-        # Forward through the underlying model
+        model_t = self._encode_time(t, x, c_noise=c_noise)  # (B,)
+        if extra_times is not None:
+            model_t = (model_t, extra_times.apply(partial(self._encode_time, x=x)))
+
         if condition is not None:
             F_x = self.model(
                 c_in * x,
-                c_noise.flatten(),
+                model_t,
                 condition=condition,
                 **model_kwargs,
             )
         else:
             F_x = self.model(
                 c_in * x,
-                c_noise.flatten(),
+                model_t,
                 **model_kwargs,
             )
 
-        D_x = c_skip * x + c_out * F_x
-
-        return D_x
+        if isinstance(F_x, tuple):
+            prediction, extra_outputs = F_x
+            return c_skip * x + c_out * prediction, extra_outputs
+        return c_skip * x + c_out * F_x
 
 
 class VPPreconditioner(BaseAffinePreconditioner):
@@ -604,6 +653,10 @@ class VPPreconditioner(BaseAffinePreconditioner):
 
     Examples
     --------
+    The VP coefficients map diffusion time to a noise level before calling
+    this convolutional backbone. The resulting prediction retains the image
+    shape for comparison with a clean-data target:
+
     >>> import torch
     >>> from physicsnemo.core import Module
     >>> # Define a simple model satisfying the diffusion model interface
@@ -742,6 +795,10 @@ class VEPreconditioner(BaseAffinePreconditioner):
 
     Examples
     --------
+    This VE preconditioner wraps a convolutional backbone whose prediction
+    matches the image shape. Passing one time per image produces one
+    preconditioned prediction per image:
+
     >>> import torch
     >>> from physicsnemo.core import Module
     >>> # Define a simple model satisfying the diffusion model interface
@@ -857,6 +914,10 @@ class IDDPMPreconditioner(BaseAffinePreconditioner):
 
     Examples
     --------
+    This iDDPM preconditioner receives a time for each image and calls a
+    convolutional backbone. It returns a clean-data prediction with the
+    same batch and spatial dimensions:
+
     >>> import torch
     >>> from physicsnemo.core import Module
     >>> # Define a simple model satisfying the diffusion model interface
@@ -1005,7 +1066,9 @@ class EDMPreconditioner(BaseAffinePreconditioner):
 
     Examples
     --------
-    Scalar ``sigma_data`` (default):
+    The first example uses one ``sigma_data`` value for all channels. A
+    convolutional backbone supplies the raw prediction, and EDM applies its
+    coefficients to produce a clean-data estimate:
 
     >>> import torch
     >>> from physicsnemo.core import Module
@@ -1024,8 +1087,8 @@ class EDMPreconditioner(BaseAffinePreconditioner):
     >>> out.shape
     torch.Size([2, 3, 16, 16])
 
-    Per-channel ``sigma_data`` for heterogeneous channels (e.g. weather
-    variables with different scales):
+    The second example gives each channel its own ``sigma_data``. This is
+    useful for variables with different scales, such as weather fields:
 
     >>> precond_ch = EDMPreconditioner(model, sigma_data=[0.3, 0.5, 0.7])
     >>> out_ch = precond_ch(x, t, condition=None)
diff --git a/physicsnemo/diffusion/utils/model_wrappers.py b/physicsnemo/diffusion/utils/model_wrappers.py
index 9837f54d3c..bed1f91b4b 100644
--- a/physicsnemo/diffusion/utils/model_wrappers.py
+++ b/physicsnemo/diffusion/utils/model_wrappers.py
@@ -18,10 +18,11 @@
 from typing import Any
 
 import torch
-from jaxtyping import Float
+from jaxtyping import Float, Shaped
 from tensordict import TensorDict
 
 from physicsnemo.core import Module
+from physicsnemo.diffusion.base import DiffusionExtraOutputs, DiffusionExtraTimes
 from physicsnemo.models.diffusion_unets import (
     DhariwalUNet,
     SongUNet,
@@ -88,8 +89,9 @@ class ConcatConditionWrapper(Module):
     x : torch.Tensor
         Noisy latent state of shape :math:`(B, *)` where :math:`B` is the
         batch size.
-    t : torch.Tensor
-        Diffusion time tensor of shape :math:`(B,)`.
+    t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes]
+        Diffusion time tensor of shape :math:`(B,)`, optionally paired with
+        extra times. The wrapped model must accept the supplied format.
     condition : torch.Tensor, TensorDict, or None, optional, default=None
         Conditioning data. Use a ``TensorDict`` to explicitly supply
         concatenated conditioning and vector conditioning as keys
@@ -101,13 +103,15 @@ class ConcatConditionWrapper(Module):
 
     Outputs
     -------
-    torch.Tensor
-        Model output with the same shape as ``x``.
+    torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs]
+        The wrapped model's prediction, optionally paired with auxiliary
+        outputs. The wrapper returns both unchanged.
 
     Examples
     --------
     Wrap :class:`~physicsnemo.models.diffusion_unets.SongUNet` with channel
-    concatenation and vector conditioning:
+    concatenation and vector conditioning. The wrapper routes the image to
+    the input channels and the vector to class labels:
 
     >>> import torch
     >>> from tensordict import TensorDict
@@ -127,7 +131,8 @@ class ConcatConditionWrapper(Module):
     >>> out.shape
     torch.Size([2, 3, 8, 8])
 
-    Wrap :class:`~physicsnemo.models.dit.DiT` similarly:
+    The same condition layout works for :class:`~physicsnemo.models.dit.DiT`.
+    The wrapper sends the vector to DiT's ``condition`` input:
 
     >>> from physicsnemo.models.dit import DiT
     >>> dit = DiT(in_channels=4, out_channels=3, condition_dim=4, input_size=8, patch_size=4)
@@ -135,6 +140,27 @@ class ConcatConditionWrapper(Module):
     >>> out = dit_wrapper(x, t, condition)
     >>> out.shape
     torch.Size([2, 3, 8, 8])
+
+    A backbone that supports extra times and outputs can use tuples. The
+    wrapper forwards these values and returns both output components:
+
+    >>> from physicsnemo.diffusion import DiffusionExtraTimes, DiffusionExtraOutputs
+    >>> class TwoTimeModel(Module):
+    ...     def forward(self, x, t, condition=None):
+    ...         t, extra_times = t
+    ...         # The first three channels exclude concatenated conditioning.
+    ...         prediction = x[:, :3] * (t - extra_times.r)[:, None, None, None]
+    ...         # Keep the required prediction separate from auxiliary outputs.
+    ...         return prediction, DiffusionExtraOutputs(
+    ...             logvar=torch.zeros_like(t), batch_size=[x.shape[0]]
+    ...         )
+    >>> wrapper = ConcatConditionWrapper(TwoTimeModel())
+    >>> times = (t, DiffusionExtraTimes(r=torch.zeros_like(t), batch_size=[2]))
+    >>> prediction, extra_outputs = wrapper(x, times, condition)
+    >>> prediction.shape, extra_outputs.logvar.shape
+    (torch.Size([2, 3, 8, 8]), torch.Size([2]))
+
+    SongUNet and DhariwalUNet backbones still require a single time tensor.
     """
 
     def __init__(
@@ -151,10 +177,14 @@ def __init__(
     def forward(
         self,
         x: Float[torch.Tensor, " B *dims"],
-        t: Float[torch.Tensor, " B"],
+        t: Float[torch.Tensor, " B"]
+        | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]],
         condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None,
         **model_kwargs: Any,
-    ) -> Float[torch.Tensor, " B *dims"]:
+    ) -> (
+        Float[torch.Tensor, " B *dims"]
+        | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]]
+    ):
         r"""
         Forward pass for the conditioned wrapper.
 
@@ -163,8 +193,9 @@ def forward(
         x : torch.Tensor
             Noisy latent state of shape :math:`(B, *)` where :math:`B` is the
             batch size.
-        t : torch.Tensor
-            Diffusion time tensor of shape :math:`(B,)`.
+        t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes]
+            Diffusion time tensor of shape :math:`(B,)`, optionally paired
+            with extra times. Forwarded to the backbone unchanged.
         condition : torch.Tensor, TensorDict, or None, optional, default=None
             Conditioning data. Use a ``TensorDict`` to supply ``cond_concat`` and
             ``cond_vec`` tensors. A plain tensor is treated as the concatenated
@@ -174,9 +205,11 @@ def forward(
 
         Returns
         -------
-        torch.Tensor
-            Model output with the same shape as ``x``.
+        torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs]
+            The wrapped model's prediction and any auxiliary outputs,
+            returned unchanged.
         """
+        t, extra_times = t if isinstance(t, tuple) else (t, None)
         cond_concat: torch.Tensor | None = None
         cond_vec: torch.Tensor | None = None
 
@@ -246,6 +279,7 @@ def forward(
         if cond_concat is not None:
             x = torch.cat([x, cond_concat], dim=1)
 
+        t = t if extra_times is None else (t, extra_times)
         if isinstance(
             self.model, (SongUNet, SongUNetPosEmbd, SongUNetPosLtEmbd, DhariwalUNet)
         ):
diff --git a/test/diffusion/helpers.py b/test/diffusion/helpers.py
index 8553d34b4e..9f96aef56a 100644
--- a/test/diffusion/helpers.py
+++ b/test/diffusion/helpers.py
@@ -20,10 +20,15 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 
 import torch
+from jaxtyping import Float, Shaped
 from tensordict import TensorDict
 
 import physicsnemo.core
 from physicsnemo.core import Module
+from physicsnemo.diffusion import (
+    DiffusionExtraOutputs,
+    DiffusionExtraTimes,
+)
 
 # Directory for test reference data
 DATA_DIR = Path(__file__).parent / "data"
@@ -77,6 +82,49 @@ def forward(self, x: torch.Tensor, t: torch.Tensor, **kwargs: Any) -> torch.Tens
         return self.net(x) / (1 + t_bc)
 
 
+class CustomExtraTimes(DiffusionExtraTimes):
+    r"""Additional time fields that callers can supply independently."""
+
+    r1: Float[torch.Tensor, " B"] | None = None
+    r2: Float[torch.Tensor, " B"] | None = None
+
+
+class StructuredDiffusionModel(Module):
+    r"""Linear model with optional destination times and learned loss weights."""
+
+    def __init__(self, output: str = "tuple"):
+        super().__init__()
+        self.output = output
+        self.gain = torch.nn.Parameter(torch.tensor(0.75))
+
+    def forward(
+        self,
+        x: Float[torch.Tensor, " B *dims"],
+        t: Float[torch.Tensor, " B"]
+        | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]],
+        condition: TensorDict | None = None,
+        *,
+        offset: float = 0.0,
+    ) -> (
+        Float[torch.Tensor, " B *dims"]
+        | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]]
+    ):
+        r"""Return a prediction, optionally paired with its loss weight."""
+        t, extra_times = t if isinstance(t, tuple) else (t, None)
+        t_bc = t.reshape(-1, *([1] * (x.ndim - 1)))  # (B, 1, ...)
+        prediction = self.gain * x + t_bc + offset
+        if extra_times is not None:
+            for time in extra_times.to_tensordict().values():
+                prediction = prediction + time.reshape(-1, *([1] * (x.ndim - 1)))
+        if condition is not None:
+            prediction = prediction + condition["shift"]
+        if self.output == "tensor":
+            return prediction
+        return prediction, DiffusionExtraOutputs(
+            logvar=self.gain * t, batch_size=[x.shape[0]]
+        )
+
+
 def make_input(
     shape: Tuple[int, ...],
     seed: int = 42,
diff --git a/test/diffusion/test_model_io.py b/test/diffusion/test_model_io.py
new file mode 100644
index 0000000000..4e4681b124
--- /dev/null
+++ b/test/diffusion/test_model_io.py
@@ -0,0 +1,208 @@
+# 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"""Check primary tensors and extra fields through diffusion wrappers."""
+
+import pytest
+import torch
+from jaxtyping import Float
+
+from physicsnemo.diffusion import DiffusionExtraOutputs, DiffusionExtraTimes
+from physicsnemo.diffusion.preconditioners import (
+    EDMPreconditioner,
+    IDDPMPreconditioner,
+    VEPreconditioner,
+    VPPreconditioner,
+)
+from physicsnemo.diffusion.utils.model_wrappers import ConcatConditionWrapper
+
+from .helpers import CustomExtraTimes, StructuredDiffusionModel
+
+CONFIGS = (
+    (EDMPreconditioner, {}),
+    (VEPreconditioner, {}),
+    (VPPreconditioner, {}),
+    (IDDPMPreconditioner, {"M": 32}),
+)
+IDS = ("edm", "ve", "vp", "iddpm")
+
+TIME_FIELDS = (None, (), ("r",), ("r1",), ("r2",), ("r1", "r2"))
+TIME_IDS = ("tensor", "empty", "r", "r1", "r2", "r1_r2")
+
+
+def _make_time(
+    t: Float[torch.Tensor, " B"], fields: tuple[str, ...] | None
+) -> Float[torch.Tensor, " B"] | tuple[Float[torch.Tensor, " B"], CustomExtraTimes]:
+    r"""Pair the primary time with the selected extra fields."""
+    if fields is None:
+        return t
+    return t, CustomExtraTimes(
+        **{name: t / (index + 2) for index, name in enumerate(fields)},
+        batch_size=t.shape,
+    )
+
+
+@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS)
+class TestConsistency:
+    @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS)
+    def test_tuple_protocol(self, preconditioner, kwargs, device, fields):
+        model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device)
+        x = torch.randn(2, 3, 5, 7, device=device)
+        t = torch.tensor([0.4, 0.8], device=device)
+        time = _make_time(t, fields)
+        calls = []
+        handle = model.model.register_forward_hook(
+            lambda module, inputs, output: calls.append((inputs, output))
+        )
+        prediction, outputs = model(x, time, offset=0.2)
+        handle.remove()
+        inputs, raw = calls[0]
+        sigma = model.sigma(t).reshape(2, 1, 1, 1)  # (B, 1, 1, 1)
+        c_in, c_noise, c_out, c_skip = model.compute_coefficients(sigma)
+        torch.testing.assert_close(inputs[0], c_in * x)
+        encoded_t = inputs[1][0] if fields is not None else inputs[1]
+        torch.testing.assert_close(encoded_t, c_noise.flatten())
+        if fields is not None:
+            assert isinstance(inputs[1][1], CustomExtraTimes)
+            for index, name in enumerate(fields):
+                original = t / (index + 2)
+                sigma_extra = model.sigma(original).reshape(2, 1, 1, 1)
+                encoded = model.compute_coefficients(sigma_extra)[1].flatten()
+                torch.testing.assert_close(getattr(inputs[1][1], name), encoded)
+                torch.testing.assert_close(getattr(time[1], name), original)
+            for name in {"r", "r1", "r2"}.difference(fields):
+                assert getattr(inputs[1][1], name) is None
+        torch.testing.assert_close(prediction, c_skip * x + c_out * raw[0])
+        assert outputs is raw[1]
+        torch.testing.assert_close(outputs.logvar, model.model.gain * encoded_t)
+
+    def test_tensor_output(self, preconditioner, kwargs, device):
+        model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device)
+        plain = preconditioner(StructuredDiffusionModel("tensor"), **kwargs).to(device)
+        x = torch.randn(2, 3, device=device)
+        t = torch.tensor([0.4, 0.8], device=device)
+        prediction, _ = model(x, t)
+        torch.testing.assert_close(prediction, plain(x, t))
+
+
+class TestConstructor:
+    def test_optional_outputs(self):
+        extras = DiffusionExtraOutputs(batch_size=[2])
+        assert extras.logvar is None
+
+    def test_conditioning_wrapper(self):
+        model = StructuredDiffusionModel()
+        wrapper = ConcatConditionWrapper(model)
+        x = torch.randn(2, 3, 5, 7)
+        t = torch.tensor([0.4, 0.8])
+        time = (t, DiffusionExtraTimes(r=t / 2, batch_size=[2]))
+        condition = torch.ones(2, 1, 5, 7)
+        prediction, extras = wrapper(x, time, condition)
+        expected, expected_extras = model(torch.cat((x, condition), dim=1), time)
+        torch.testing.assert_close(prediction, expected)
+        torch.testing.assert_close(extras.logvar, expected_extras.logvar)
+
+
+@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS)
+class TestGradientFlow:
+    @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS)
+    def test_numerical_gradients(self, preconditioner, kwargs, fields):
+        model = preconditioner(StructuredDiffusionModel(), **kwargs).double()
+        x = torch.randn(2, 3, dtype=torch.float64, requires_grad=True)
+        t = torch.tensor([0.4, 0.8], dtype=torch.float64, requires_grad=True)
+
+        def forward(x, t):
+            prediction, extras = model(x, _make_time(t, fields))
+            return prediction, extras.logvar
+
+        assert torch.autograd.gradcheck(forward, (x, t))
+
+
+@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS)
+class TestCompile:
+    @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS)
+    def test_forward_backward(self, preconditioner, kwargs, device, fields):
+        model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device)
+        x = torch.randn(2, 3, 5, 7, device=device, requires_grad=True)
+        t = torch.tensor([0.4, 0.8], device=device, requires_grad=True)
+        time = _make_time(t, fields)
+        expected, expected_extras = model(x, time)
+        expected_grads = torch.autograd.grad(
+            expected.sum() + expected_extras.logvar.sum(),
+            (x, t, model.model.gain),
+            retain_graph=True,
+        )
+        compiled = torch.compile(model, backend="aot_eager", fullgraph=True)
+        result, extras = compiled(x, time)
+        with torch._dynamo.config.patch(error_on_recompile=True):
+            repeated, _ = compiled(x, time)
+        torch.testing.assert_close(repeated, result)
+        torch.testing.assert_close(result, expected)
+        torch.testing.assert_close(extras.logvar, expected_extras.logvar)
+        gradients = torch.autograd.grad(
+            result.sum() + extras.logvar.sum(), (x, t, model.model.gain)
+        )
+        for actual, reference in zip(gradients, expected_grads):
+            torch.testing.assert_close(actual, reference)
+
+    def test_conditioning_wrapper(self, preconditioner, kwargs, device):
+        wrapper = ConcatConditionWrapper(
+            preconditioner(StructuredDiffusionModel(), **kwargs)
+        ).to(device)
+        x = torch.randn(2, 3, device=device, requires_grad=True)
+        t = torch.tensor([0.4, 0.8], device=device)
+        time = (t, DiffusionExtraTimes(r=t / 2, batch_size=[2]))
+        expected, _ = wrapper(x, time)
+        compiled = torch.compile(wrapper, backend="aot_eager", fullgraph=True)
+        actual, _ = compiled(x, time)
+        torch.testing.assert_close(actual, expected)
+        torch.testing.assert_close(
+            torch.autograd.grad(actual.sum(), x)[0],
+            torch.autograd.grad(expected.sum(), x)[0],
+        )
+
+
+class TestDiagnostics:
+    @pytest.mark.parametrize(
+        "t,r,batch_size",
+        (
+            (torch.ones(2, 1), torch.ones(2), [2]),
+            (torch.ones(2), torch.ones(2), []),
+        ),
+    )
+    def test_invalid_time(self, t, r, batch_size):
+        model = EDMPreconditioner(StructuredDiffusionModel())
+        with pytest.raises(ValueError, match="shape|batch"):
+            model(
+                torch.ones(2, 3), (t, DiffusionExtraTimes(r=r, batch_size=batch_size))
+            )
+
+
+@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS)
+class TestNonRegression:
+    @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS)
+    def test_checkpoint(self, preconditioner, kwargs, tmp_path, fields):
+        model = preconditioner(StructuredDiffusionModel(), **kwargs)
+        x = torch.randn(2, 3)
+        times = _make_time(torch.tensor([0.4, 0.8]), fields)
+        expected, expected_extras = model(x, times)
+        checkpoint = tmp_path / "model.pt"
+        torch.save(model.state_dict(), checkpoint)
+        restored = preconditioner(StructuredDiffusionModel(), **kwargs)
+        restored.load_state_dict(torch.load(checkpoint, weights_only=True))
+        actual, extras = restored(x, times)
+        torch.testing.assert_close(actual, expected)
+        torch.testing.assert_close(extras.logvar, expected_extras.logvar)