feat(training): centralize WGAN-GP + diffusion consumers through fused primitives (#1847)

Closes #1845 and #1846. Routes all four WGAN-GP critics through WganGpFusedStep
(the fused-plan primitive from PR #1843) and adds MultiSlotFusedStep wire-ups
to the three tabular diffusion consumers plus a base-class opt-in hook for
DiffusionModelBase subclasses.

## Optimizer config plumbing

Zero new API surface: IFusedOptimizerSpec.TryGetFusedOptimizerConfig already
exposes (OptimizerType, LR, Beta1, Beta2, Epsilon, WeightDecay, Schedule,
UseBf16Moments). Widened NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig
from private to internal so the sibling generators in the same assembly can
reuse the existing helper.

## WGAN-GP consumers (#1845)

* CTGANGenerator, CopulaGANGenerator, TableGANGenerator, CausalGANGenerator —
  each critic training method now attempts WganGpFusedStep.TryStep FIRST with
  the discriminator's optimizer hyperparameters extracted via
  TryMapToFusedOptimizerConfig, falls back to the existing
  GpuResidentFusedStep path (secondary fused), then the eager tape (final
  fallback). The ε ∈ [0, 1]^B epsilon sampler uses
  Engine.TensorRandomUniformRange to match each critic's local
  ComputeGradientPenalty behavior.
* Non-Adam optimizers (Lion, LBFGS) that don't implement IFusedOptimizerSpec
  cleanly fall through — TryMapToFusedOptimizerConfig returns false and the
  code path skips to GpuResidentFusedStep as before.

## Diffusion consumers (#1846)

* TabDDPMGenerator — refactored to expose a slot-based forward
  (BuildTabDDPMSlots + DenoiserForwardFromTensors +
  ComputeDiffusionLossTapeFromTensors). Per-row TrainBatch loop now attempts
  MultiSlotFusedStep with (numNoisy, actualNoise, catNoisy, catClean,
  rawSinusoidalTimeEmbed) as persistent slots. The learnable
  _timestepProjection stays INSIDE the compiled forward closure so its
  weights participate in the backward pass. Plan is compiled once on the
  first row and replayed via slot-data refresh for subsequent rows.
* TabSynGenerator — TrainDiffusionBatch's per-row loop wired with
  MultiSlotFusedStep on (noisyLatent, actualNoise, projectedTimeEmbed).
  Matches the existing eager path's semantic that _timestepProjection is
  NOT in _diffMLPLayers (kept detached in the eager path too), so the
  projected embedding is precomputed host-side per row and passed as slot
  data.
* Finance/Forecasting/Foundation/CSDI —
  - ApplyInstanceNormalization rewritten with traceable engine ops
    (ReduceMean + ReduceVariance + TensorSqrt + broadcast subtract/divide) —
    same pattern as the TFC RevIN fix. The previous `.Data.Span` per-batch
    loop froze at trace time.
  - New BuildCsdiSlots + DenoiserForwardFromSlots express the DDPM x_t
    formation and packed denoising input via TensorConcatenate + engine
    scalar multiplies. Replaces the `.Data.Span[i] = xt[0, i]` fill that
    baked the trace batch's x_t into the compiled plan.
  - Train() attempts MultiSlotFusedStep first, falls back to the existing
    eager ComputeDenoisingPairTape path when the fused path can't engage.
* Diffusion/DiffusionModelBase —
  - New opt-in `protected virtual bool SupportsFusedDenoising => false;`
    property. Base default is false so no existing subclass changes
    behavior.
  - Train() attempts MultiSlotFusedStep when SupportsFusedDenoising is true
    AND the training optimizer maps cleanly to a fused config. Slots:
    (noisySample, noise). Loss = MSE(pred, noise). QAT shadow restoration
    is preserved on the fused-success path.
  - Subclasses with fully-traceable PredictNoise / PredictNoiseBatched
    (e.g. after auditing to remove `.Data.Span` host loops) can opt in via
    a single-line override; no infrastructure changes needed elsewhere.

## Verification

* net8.0, net471, net10.0 all build clean.
* No API surface changes on IGradientBasedOptimizer<T> — the existing
  IFusedOptimizerSpec interface (already implemented by all fuse-able
  optimizers) provided everything needed.
* All consumers preserve eager fallback path for non-fuse-able optimizers
  and non-GPU hosts.

Co-authored-by: franklinic <franklin@ivorycloud.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

Author: 3 people

src/Diffusion/DiffusionModelBase.cs

@@ -126,6 +126,29 @@ protected virtual IEnumerable<IDisposable> EnumerateDisposableComponents()
     /// <inheritdoc/>
     public virtual ModelOptions GetOptions() => _options;
 
+    /// <summary>
+    /// Opt-in flag: does this subclass's <see cref="PredictNoise"/> /
+    /// <see cref="PredictNoiseBatched"/> path use ONLY traceable engine ops (no
+    /// host-side <c>.Data.Span</c> loops, no per-step class-field mutation)?
+    /// When <c>true</c>, <see cref="Train"/> attempts a
+    /// <see cref="AiDotNet.Training.MultiSlotFusedStep{T}"/> fused-plan step
+    /// with (noisySample, noise) as persistent slots — the compiled plan replays
+    /// the forward per training call so <c>PredictNoise</c> must be free of
+    /// per-call side effects that would freeze at trace time.
+    /// <para>
+    /// Base default is <c>false</c> so existing subclasses (which may still use
+    /// <c>.Data.Span</c> internally) run through the eager tape as before. Override
+    /// after auditing your forward path. See ooples/AiDotNet#1846.
+    /// </para>
+    /// </summary>
+    protected virtual bool SupportsFusedDenoising => false;
+
+    private static void RestoreShadow(Tensor<T> param, Vector<T> shadow)
+    {
+        var span = param.Data.Span;
+        for (int k = 0; k < span.Length && k < shadow.Length; k++) span[k] = shadow[k];
+    }
+
     /// <summary>
     /// The optional neural network architecture blueprint for custom layer configuration.
     /// </summary>
@@ -1132,6 +1155,70 @@ public virtual void Train(Tensor<T> input, Tensor<T> expectedOutput)
             noisySampleTensor = new Tensor<T>(input._shape, noisySample);
         }
 
+        // Preferred fused path: MultiSlotFusedStep with (noisySample, noise,
+        // timestepsPerElement) as persistent slots. Only engages when the
+        // concrete subclass has certified its PredictNoise/PredictNoiseBatched
+        // path as fully traceable by opting in via SupportsFusedDenoising —
+        // eager PredictNoise implementations that use host-side .Data.Span
+        // loops would freeze at trace time. See ooples/AiDotNet#1846.
+        // Only try the fused path when an optimizer has already been resolved
+        // (avoids double-construction with a slightly-different config).
+        var trainingOptimizerForFused = _trainingOptimizer;
+        if (SupportsFusedDenoising
+            && trainingOptimizerForFused is not null
+            && NeuralNetworks.NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig(
+                trainingOptimizerForFused,
+                out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2,
+                out var mfsEps, out var mfsWd, out _, out _))
+        {
+            var trainableForFused = CollectTrainableParameters();
+            var noiseSlotT = new Tensor<T>(noisySampleTensor._shape, noiseVector);
+            var slots = new Tensor<T>[]
+            {
+                noisySampleTensor,
+                noiseSlotT,
+            };
+            using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep<T>();
+            var timestepsSnapshot = timesteps;
+            var timestepSnapshotSingle = timestep;
+            var isBatchedSnapshot = isBatched;
+            Tensor<T> ForwardFromSlots(IReadOnlyList<Tensor<T>> s)
+            {
+                return isBatchedSnapshot
+                    ? PredictNoiseBatched(s[0], timestepsSnapshot)
+                    : PredictNoise(s[0], timestepSnapshotSingle);
+            }
+            Tensor<T> ComputeLossFromSlots(Tensor<T> pred, IReadOnlyList<Tensor<T>> s)
+            {
+                var diff = Engine.TensorSubtract(pred, s[1]);
+                var sq = Engine.TensorMultiply(diff, diff);
+                var axes = Enumerable.Range(0, sq.Shape.Length).ToArray();
+                return Engine.ReduceMean(sq, axes, keepDims: false);
+            }
+            if (trainableForFused.Length > 0
+                && multiSlotStep.TryStep(
+                    parameters: trainableForFused,
+                    zeroGradAction: null,
+                    freshSlotData: slots,
+                    forward: ForwardFromSlots,
+                    computeLoss: ComputeLossFromSlots,
+                    optimizerType: mfsOptType,
+                    learningRate: mfsLr,
+                    beta1: mfsB1,
+                    beta2: mfsB2,
+                    epsilon: mfsEps,
+                    weightDecay: mfsWd,
+                    out T _))
+            {
+                if (qatShadows is not null && qatParams is not null)
+                {
+                    for (int i = 0; i < qatParams.Length; i++)
+                        RestoreShadow(qatParams[i], qatShadows[i]);
+                }
+                return;
+            }
+        }
+
         using var tape = new GradientTape<T>();
 
         // Forward pass — triggers lazy layer initialization, then we walk for

src/Finance/Forecasting/Foundation/CSDI.cs

@@ -287,16 +287,51 @@ public override void Train(Tensor<T> input, Tensor<T> target)
 
         var trainableParams = Training.TapeTrainingStep<T>.CollectParameters(Layers).ToArray();
 
-        // CSDI's stochastic denoising step cannot be traced through the compiled
-        // fused plan on the current Tensors build. ComputeDenoisingPairTape samples
-        // a fresh (timestep, noise) pair each call, and calling it twice — once in
-        // the forward closure, once in the loss closure — produces two independent
-        // samples that don't correspond to each other. Since the compiled plan
-        // replays the SAME captured graph every step, we can't refresh the RNG
-        // sample per replay either. This path stays on the eager tape until
-        // Tensors' PersistentInputRegistry (PR #763) lands, at which point the
-        // (timestep, noise) pair can be sampled per step and passed in as data.
+        // Preferred fused path: MultiSlotFusedStep with the sampled (noise,
+        // xt-scale, sinT) tuple passed as persistent slots. Refreshes per step
+        // by host-sampling a fresh (t, ε) pair and copying values into the
+        // slot tensors — the compiled forward reads the CURRENT slot data on
+        // every replay. See ooples/AiDotNet#1846.
+        if (trainableParams.Length > 0
+            && NeuralNetworks.NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig(
+                _optimizer,
+                out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2,
+                out var mfsEps, out var mfsWd, out _, out _))
+        {
+            var slots = BuildCsdiSlots(input, target);
+            if (slots is not null)
+            {
+                using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep<T>();
+                Tensor<T> ForwardFromSlots(IReadOnlyList<Tensor<T>> s) => DenoiserForwardFromSlots(s);
+                Tensor<T> ComputeLossFromSlots(Tensor<T> pred, IReadOnlyList<Tensor<T>> s)
+                {
+                    // s[1] = ε_true (noise slot). Loss = MSE(ε_pred, ε_true) via
+                    // the model's LossFunctionBase so custom losses are respected.
+                    return loss.ComputeTapeLoss(pred, s[1]);
+                }
+                if (multiSlotStep.TryStep(
+                        parameters: trainableParams,
+                        zeroGradAction: null,
+                        freshSlotData: slots,
+                        forward: ForwardFromSlots,
+                        computeLoss: ComputeLossFromSlots,
+                        optimizerType: mfsOptType,
+                        learningRate: mfsLr,
+                        beta1: mfsB1,
+                        beta2: mfsB2,
+                        epsilon: mfsEps,
+                        weightDecay: mfsWd,
+                        out T fusedLoss))
+                {
+                    LastLoss = fusedLoss;
+                    return;
+                }
+            }
+        }
 
+        // Eager fallback: samples (t, ε) inline and runs the denoising pair on
+        // the tape. Same objective; used when the fused path can't engage
+        // (non-fuse-able optimizer, non-GPU host, etc.).
         using var tape = new GradientTape<T>();
         var (epsilonPred, epsilonTarget) = ComputeDenoisingPairTape(input, target);
 
@@ -352,6 +387,112 @@ public override void Train(Tensor<T> input, Tensor<T> target)
         }
     }
 
+    /// <summary>
+    /// Assembles the persistent-slot data for the MultiSlotFusedStep wire-up:
+    /// samples (t, ε) host-side then packs everything the compiled forward
+    /// needs as tensor slots. Slot layout:
+    /// <list type="bullet">
+    /// <item><description>[0] target — the clean values to denoise (Ho 2020 x_0).</description></item>
+    /// <item><description>[1] ε_true — freshly-sampled Gaussian noise, same shape as target.</description></item>
+    /// <item><description>[2] sqrtAlphaBar_scalar — cumulative α scaling at the sampled t, shape [1].</description></item>
+    /// <item><description>[3] sqrtOneMinusAlphaBar_scalar — noise scaling at t, shape [1].</description></item>
+    /// <item><description>[4] sinT_scalar — sin(2π·t/T) timestep encoding, shape [1].</description></item>
+    /// <item><description>[5] conditioned — RevIN-normalized observed input, shape [1, seqLen].</description></item>
+    /// </list>
+    /// Returns null when the target is degenerate (zero length) so the caller
+    /// falls back to eager training.
+    /// </summary>
+    private IReadOnlyList<Tensor<T>>? BuildCsdiSlots(Tensor<T> input, Tensor<T> target)
+    {
+        int targetLen = target.Length;
+        if (targetLen <= 0) return null;
+
+        var rand = RandomHelper.CreateSecureRandom();
+        int t = rand.Next(_numDiffusionSteps);
+        var noiseData = new T[targetLen];
+        for (int i = 0; i < targetLen; i++) noiseData[i] = SampleStandardNormal(rand);
+        var epsilonTrue = new Tensor<T>(target._shape, new Vector<T>(noiseData));
+
+        T sqrtAlphaBar = NumOps.Sqrt(_alphasCumprod[t]);
+        T sqrtOneMinus = NumOps.Sqrt(NumOps.Subtract(NumOps.One, _alphasCumprod[t]));
+        T sinT = NumOps.FromDouble(Math.Sin(2.0 * Math.PI * t / Math.Max(1, _numDiffusionSteps - 1)));
+
+        var sqrtAlphaBarT = new Tensor<T>(new[] { 1 });
+        sqrtAlphaBarT[0] = sqrtAlphaBar;
+        var sqrtOneMinusT = new Tensor<T>(new[] { 1 });
+        sqrtOneMinusT[0] = sqrtOneMinus;
+        var sinTT = new Tensor<T>(new[] { 1 });
+        sinTT[0] = sinT;
+
+        var conditioned = ApplyInstanceNormalization(input);
+        if (conditioned.Rank == 1)
+            conditioned = Engine.Reshape(conditioned, new[] { 1, conditioned.Length });
+        return new[] { target, epsilonTrue, sqrtAlphaBarT, sqrtOneMinusT, sinTT, conditioned };
+    }
+
+    /// <summary>
+    /// Fused-path denoiser forward: reads the slot tuple from
+    /// <see cref="BuildCsdiSlots"/>, builds x_t via traceable engine
+    /// arithmetic (no <c>.Data.Span</c> pack — the previous inline
+    /// implementation froze the trace batch's x_t into the compiled plan),
+    /// concatenates [x_t | conditioning_slice | sin(t)] via
+    /// <see cref="IEngine.TensorConcatenate{T}"/>, and runs the projection +
+    /// residual stack + output projection. Returns predicted noise aligned
+    /// with the noise slot's shape.
+    /// </summary>
+    private Tensor<T> DenoiserForwardFromSlots(IReadOnlyList<Tensor<T>> slots)
+    {
+        var target = slots[0];
+        var epsilonTrue = slots[1];
+        var sqrtAlphaBarT = slots[2];
+        var sqrtOneMinusT = slots[3];
+        var sinTT = slots[4];
+        var conditioned = slots[5];
+
+        int targetLen = target.Length;
+        // x_t = sqrt(α̅_t) * target + sqrt(1-α̅_t) * ε — all traceable.
+        var scaledTarget = Engine.TensorMultiplyScalar(target, sqrtAlphaBarT[0]);
+        var scaledNoise = Engine.TensorMultiplyScalar(epsilonTrue, sqrtOneMinusT[0]);
+        var xt = Engine.TensorAdd(scaledTarget, scaledNoise);
+        var xt1d = xt.Rank == 1 ? xt : Engine.Reshape(xt, new[] { targetLen });
+
+        // Conditioning slice: first Min(condLen, hiddenDimension) elements,
+        // flattened to 1-D so we can concat with xt1d + sinT.
+        var condFlat = conditioned.Rank == 1
+            ? conditioned
+            : Engine.Reshape(conditioned, new[] { conditioned.Length });
+        int condLen = Math.Min(condFlat.Length, _hiddenDimension);
+        var condSlice = condLen == condFlat.Length
+            ? condFlat
+            : Engine.TensorSlice(condFlat, new[] { 0 }, new[] { condLen });
+
+        // Pack [xt | conditioning | sin(t)] via TensorConcatenate — replaces
+        // the .Data.Span[i] fill that froze at trace time.
+        var packed1D = Engine.TensorConcatenate(new[] { xt1d, condSlice, sinTT }, axis: 0);
+        var denoisingInput = Engine.Reshape(packed1D, new[] { 1, targetLen + condLen + 1 });
+
+        var eps = (Tensor<T>)denoisingInput;
+        if (_inputProjection is not null)
+            eps = _inputProjection.Forward(eps);
+        foreach (var layer in _residualLayers)
+            eps = layer.Forward(eps);
+        if (_outputProjection is not null)
+            eps = _outputProjection.Forward(eps);
+
+        if (eps.Rank == 2 && eps.Shape[1] > epsilonTrue.Length)
+            eps = Engine.TensorNarrow(eps, dim: 1, start: 0, length: epsilonTrue.Length);
+        if (eps.Length != epsilonTrue.Length)
+        {
+            throw new InvalidOperationException(
+                $"CSDI denoising pair: predicted-noise length ({eps.Length}, shape=["
+                + $"{string.Join(",", eps._shape)}]) does not match true-noise length ("
+                + $"{epsilonTrue.Length}, shape=[{string.Join(",", epsilonTrue._shape)}]).");
+        }
+        if (!eps._shape.AsEnumerable().SequenceEqual(epsilonTrue._shape))
+            eps = Engine.Reshape(eps, epsilonTrue._shape);
+        return eps;
+    }
+
     /// <summary>
     /// Builds the (predicted-noise, true-noise) pair for one DDPM
     /// training step. Samples a timestep and a fresh noise tensor,
@@ -535,23 +676,28 @@ public override Dictionary<string, T> Evaluate(Tensor<T> predictions, Tensor<T>
         return new Dictionary<string, T> { ["MSE"] = mse, ["MAE"] = mae, ["RMSE"] = NumOps.Sqrt(mse) };
     }
 
+    /// <inheritdoc/>
+    /// <remarks>
+    /// Traceable RevIN forward — same pattern as
+    /// <see cref="TFC{T}.ApplyInstanceNormalization"/>. Uses
+    /// <see cref="IEngine.ReduceMean{T}"/> / <see cref="IEngine.ReduceVariance{T}"/>
+    /// / <see cref="IEngine.TensorSqrt{T}"/> / broadcast subtract+divide so the
+    /// per-batch stats re-execute on every replay under a compiled fused plan.
+    /// </remarks>
     public override Tensor<T> ApplyInstanceNormalization(Tensor<T> input)
     {
         int batchSize = input.Rank > 1 ? input.Shape[0] : 1;
         int seqLen = input.Rank > 1 ? input.Shape[1] : input.Length;
-        var result = new Tensor<T>(input._shape);
-        for (int b = 0; b < batchSize; b++)
-        {
-            T mean = NumOps.Zero;
-            for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) mean = NumOps.Add(mean, input[idx]); }
-            mean = NumOps.Divide(mean, NumOps.FromDouble(seqLen));
-            T variance = NumOps.Zero;
-            for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) { var diff = NumOps.Subtract(input[idx], mean); variance = NumOps.Add(variance, NumOps.Multiply(diff, diff)); } }
-            variance = NumOps.Divide(variance, NumOps.FromDouble(seqLen));
-            T std = NumOps.Sqrt(NumOps.Add(variance, NumOps.FromDouble(1e-5)));
-            for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length && idx < result.Length) result.Data.Span[idx] = NumOps.Divide(NumOps.Subtract(input[idx], mean), std); }
-        }
-        return result;
+        if (seqLen <= 0) return input;
+
+        bool reshaped = input.Rank != 2;
+        var flat = reshaped ? Engine.Reshape(input, new[] { batchSize, seqLen }) : input;
+        var mean = Engine.ReduceMean(flat, new[] { 1 }, keepDims: true);
+        var variance = Engine.ReduceVariance(flat, new[] { 1 }, keepDims: true);
+        var std = Engine.TensorSqrt(Engine.TensorAddScalar(variance, NumOps.FromDouble(1e-5)));
+        var centered = Engine.TensorBroadcastSubtract(flat, mean);
+        var normalized = Engine.TensorBroadcastDivide(centered, std);
+        return reshaped ? Engine.Reshape(normalized, input._shape) : normalized;
     }
 
     public override Dictionary<string, T> GetFinancialMetrics()

src/NeuralNetworks/NeuralNetworkBase.cs

@@ -8800,7 +8800,7 @@ private static void EmitFusedPathEventIfEnabled(bool hit, string? reason)
     /// configure-once contract: adaptive learning rates, attached LR schedulers,
     /// or AMSGrad mode (which the fused kernel doesn't model).
     /// </summary>
-    private static bool TryMapToFusedOptimizerConfig(
+    internal static bool TryMapToFusedOptimizerConfig(
         IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>> optimizer,
         out AiDotNet.Tensors.Engines.Compilation.OptimizerType optimizerType,
         out float learningRate,

src/NeuralNetworks/SyntheticData/CTGANGenerator.cs

@@ -614,7 +614,46 @@ private void TrainDiscriminatorStepBatched(Matrix<T> transformedData, int numPac
 
         var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks);
 
-        // GPU-RESIDENT WGAN-GP disc via Tensors PR #763.
+        // Preferred fused path: WganGpFusedStep runs the full WGAN-GP objective
+        // (Wasserstein + λ·GP with createGraph=true GP) in one compiled plan,
+        // with persistent (real, fake, ε) slots refreshed each Step. Falls
+        // through to the GpuResidentFusedStep path (which uses this critic's
+        // per-batch epsilon sampled inline) when the optimizer has no fused-
+        // kernel mapping (adaptive LR, non-fuse-able type) or the primitive
+        // itself declines. See ooples/AiDotNet#1845.
+        var discParams = TapeTrainingStep<T>.CollectParameters(_discLayers);
+        if (discParams.Count > 0
+            && NeuralNetworks.NeuralNetworkBase<T>.TryMapToFusedOptimizerConfig(
+                _discriminatorOptimizer,
+                out var wganOptType, out var wganLr, out var wganB1,
+                out var wganB2, out var wganEps, out var wganWd,
+                out _, out _))
+        {
+            using var wganStep = new AiDotNet.Training.WganGpFusedStep<T>();
+            Tensor<T> DiscFwd(Tensor<T> inp) => DiscriminatorForwardBatched(inp, isTraining: true);
+            Tensor<T> EpsilonSampler(int bs) =>
+                Engine.TensorRandomUniformRange<T>(new[] { bs, 1 }, NumOps.Zero, NumOps.One);
+            if (wganStep.TryStep(
+                    discParameters: discParams,
+                    realBatch: realPacked,
+                    fakeBatch: fakePacked,
+                    discForward: DiscFwd,
+                    epsilonSampler: EpsilonSampler,
+                    gradientPenaltyWeight: _options.GradientPenaltyWeight,
+                    optimizerType: wganOptType,
+                    learningRate: wganLr,
+                    beta1: wganB1,
+                    beta2: wganB2,
+                    epsilon: wganEps,
+                    weightDecay: wganWd,
+                    out T _))
+            {
+                return;
+            }
+        }
+
+        // Secondary fused path: GpuResidentFusedStep with the loss composed via
+        // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844).
         var trainableDiscLayers = _discLayers.OfType<ITrainableLayer<T>>().ToList();
         if (trainableDiscLayers.Count > 0)
         {
@@ -650,7 +689,8 @@ Tensor<T> Loss(Tensor<T> allScores, Tensor<T> _)
         }
 
         using var tape = new GradientTape<T>();
-        var discParams = TapeTrainingStep<T>.CollectParameters(_discLayers);
+        // discParams already collected above for the WganGpFusedStep attempt;
+        // reuse it here to avoid a redundant Layers → parameter scan.
 
         var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true);
         var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true);