feat(training): GPU-resident fused step for non-TS single-net models (#1843)

* feat(training): gPU-resident fused step for non-TS single-net models

Kicks off the non-time-series GPU-residency sweep. Adds a shared helper for
classes that don't inherit from NeuralNetworkBase or TimeSeriesModelBase but
still want their Train() to route forward + backward + optimizer through the
compiled fused plan (weights / activations / Adam moments resident on-device
across the whole step).

* src/Training/GpuResidentFusedStep.cs — shared helper. TryResolveOptimizerConfig
  maps a runtime IGradientBasedOptimizer to the fused-plan OptimizerType +
  hyperparameters (Adam / AdamW / SGD supported via case-insensitive class-name
  match + reflection over Options.InitialLearningRate / Beta1 / Beta2 / Epsilon
  / WeightDecay). TryStep is the one-shot entry.

Wired the following single-net models through the helper (mirrors
NeuralNetworkBase.TrainWithFusedStep):

* GraphClassificationModel (NN base + GNN + pooling + cross-entropy) —
  routes tape training through the fused plan when float + DirectGpu +
  compilation are live; falls through to the existing eager tape+optimizer
  loop on any failure.
* LinkPredictionModel (NN base + GNN + node embeddings + BCE) — same
  wire-up as GraphClassificationModel.
* FourierNeuralOperator (lift → FourierLayers → project, hardcoded SGD +
  learning rate 0.001) — captures the whole spectral-conv chain in a fused
  SGD plan; falls through to the in-place SGD loop below.

GAN-family generators (CTGAN, DPCTGAN, PATEGAN, ...), CLAPModel (learned
scalar outside layer hierarchy), AutoDiffTabGenerator (per-sample variable
shape defeats plan replay) and other non-conforming shapes are queued for
follow-up commits. The shared helper is the common seam so they all wire
the same way once their per-step structure fits the fused-plan contract
(constant shape + Adam/AdamW/SGD optimizer + ITrainableLayer-carried params).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): gPU-resident fused step for Finance foundation forecasters

Wires TFC, TOTEM, CSDI, MQCNN through the compiled fused SGD plan (all four use
in-place SGD with lr=0.001 as their eager path). Extends the shared helper with
an inline IsGpuResidentAvailable check so callers don't need to duplicate the
availability gate.

* GpuResidentFusedStep — adds IsGpuResidentAvailable static property (mirrors
  TimeSeriesModelBase.CanTrainOnGpu / NeuralNetworkBase.CanTrainOnGpu). TryStep
  short-circuits when unavailable so callers stay on their eager path.
* TFC — supervised + contrastive branches fused into one closure; the fused
  plan captures both losses in a single backward.
* TOTEM — reconstruction + VQ commitment terms fused; the recompute-loss
  closure re-derives the commitment loss on each replay for graph parity.
* CSDI — denoising-score-matching: the forward closure re-samples timestep +
  noise each step (via ComputeDenoisingPairTape) so replay produces fresh
  training data even though the plan's captured graph shape is fixed.
* MQCNN — multi-quantile pinball loss captured directly via ComputeMultiQuantilePinballLossTape.

All four fall through to the existing in-place SGD path on any fused-plan failure
(unsupported op, non-compilable graph, etc). Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): gPU-resident fused generator step for GAN-family synthetic generators + TVAE

Wires the generator side of the dual-network training loop through the fused
compiled plan for 9 SyntheticData generators. Each covers a distinct pattern:

* PATEGAN — per-sample loop; fused plan compiles on first sample, replays
  across the batch with fresh noise per iteration.
* MedSynth — batched noise -> decoder -> discriminator-frozen -> non-saturating
  log-sigmoid loss.
* CausalGAN — batched noise -> generator -> optional causal-structure ->
  output activations -> disc-frozen -> -avgFake.
* OCTGAN — per-sample loop; noise -> gen -> disc-frozen embedding -> SVDD dist².
* TableGAN — noise + real batch -> gen -> composite loss (fake-scores +
  information-loss + optional classification).
* CTGAN — the tricky one: pack (noise, cond, mask) into one persistent input
  tensor so the closure can slice cond and mask back out on replay; captures
  both -avgFake AND conditional cross-entropy on the fused plan.
* DPCTGAN — same pattern as CTGAN but simpler (no mask, no CE).
* CopulaGAN — same pattern as DPCTGAN.
* TVAE — encoder + reparam + decoder + composite ELBO (recon + KL) all fused;
  Reparameterize re-samples inside the closure so training stays stochastic.

Discriminator/critic/student layers are NOT passed to the fused step (their
weights stay frozen on the gen step, matching the eager path semantics).
Falls through to the eager tape+optimizer path on any failure.

The discriminator STEPS of these GANs are not yet wired — they need a separate
pass because their loss involves gradient penalty (WGAN-GP) or per-example DP
clipping (DPCTGAN) which don't compose cleanly with the single-closure fused
step yet.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): gPU-resident fused step for TabSyn VAE + MisGAN mask/data gens

Two more SyntheticData generators covered — each covers a distinct pattern:

* TabSyn.TrainVAEBatch — per-row VAE ELBO training. Fused plan compiles on
  first row, replays across the batch with refreshed input per iteration.
  Reparameterize re-samples inside the closure so training stays stochastic.
* MisGAN.TrainDataGeneratorStep — dual-noise generator (data noise + mask
  noise) packed into a single persistent input tensor; closure slices them
  back out on replay. Loss = -E[D_x(fakeRow ⊙ fakeMask)].
* MisGAN.TrainMaskGeneratorStep — simpler single-noise version; per-sample
  loop with fused plan replay across the batch.

The DataDiscriminator step (WGAN-style critic + weight clipping) and mask
discriminator step aren't wired — they need separate plans and the
ClipWeights post-step interacts poorly with the fused optimizer's
in-place update. Follow-up.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): gPU-resident fused step for TimeGAN P1/P2 + MedSynth non-DP disc

* TimeGAN Phase 1 (TrainReconstructionStepBatched) — embedder + recovery
  reconstruction MSE. Straight single-network fused-step conversion.
* TimeGAN Phase 2 (TrainSupervisedStepBatched) — supervisor's next-step MSE
  in the embedded space, with the embedder frozen (its layers not in the
  trainable set for this step). Two-tensor closure (ht, htNext).
* MedSynth non-DP disc step — pack (real, fake) into a single input tensor
  along axis 0; slice back in the loss for BCE-real + BCE-fake. Generator
  runs OUTSIDE the fused plan so its weights stay frozen on the critic step.

TimeGAN Phase 3 (adversarial + supervised joint) and MedSynth DP-SGD paths
still use their eager tape — the WGAN-GP gradient penalty and per-example
DP clipping don't compose with the single-closure fused optimizer step.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): extra-tensors parameter on fused-step API

Extends the fused compiled-plan training step to accept raw trainable tensors
that aren't naturally carried by an ITrainableLayer (e.g. GOGGLE's soft
adjacency A; CLAP's learned _logTemperature scalar). Solves the SOLID
concerns raised in review of the ITrainableLayer-wrapper alternative:
* ISP — no forcing raw tensors to implement the full ILayer surface with
  no-op Forward / SetTrainingMode / GetParameterGradients members
* LSP — no risk of "layer" wrappers with identity Forward diverging from
  the layer contract's behavioral expectations elsewhere in the codebase

Wire-up:
* CompiledTapeTrainingStep.TryStepWithFusedOptimizer — new optional
  extraTensors param. Threaded into the dedup-aware parameter collection
  (CollectDeduplicatedParametersWithExtras) so extras get moment buffers,
  gradient accumulation, and GPU-residency in exactly the same code path
  that layer-carried params do.
* GpuResidentFusedStep.TryStep — plumbs extras through. A callsite with
  only extras (no layers) is a valid config for models whose whole
  trainable surface is raw tensors.

Dedup is by Tensor<T> reference across BOTH sources, so an extra tensor
that also happens to be layer-carried is registered exactly once — the
same shared/tied-weight protection the layer-only collector provides.

Enables Phase 4E (GOGGLE + CLAP wire-ups).
Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): wire GOGGLE + CLAP through fused-step extra-tensors API

Consumes Phase 4A's extraTensors parameter to route models with raw
trainable state (outside the ITrainableLayer contract) through the fused
compiled plan.

* GOGGLE — soft adjacency A (initialised in InitializeModel, projected
  after each optimizer step via ProjectAdjacencyConstraints) is passed
  through extraTensors. The fused optimizer allocates its moment buffer,
  accumulates gradients, and applies the update in-place. Reparameterize
  re-samples inside the forward closure so training stays stochastic;
  ProjectAdjacencyConstraints runs after the fused step to keep the
  adjacency on the valid soft-adjacency manifold.
* CLAP — the learned _logTemperature scalar (Radford 2021 / Wu 2023
  contrastive-alignment temperature) goes through extraTensors. Both
  audio and text encoders are in the layers list; the fused step handles
  all three parameter classes uniformly.

Both fall through to the existing eager tape+optimizer path on any
failure (unsupported optimizer, non-compilable graph, etc).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(diffusion): batched per-element timesteps for DDPM-canonical training

Adds industry-standard batched-per-element timestep API to the diffusion stack
(Ho et al. 2020 §3, HuggingFace diffusers reference). Previously
DiffusionModelBase.Train sampled ONE timestep for the whole batch — reducing
the signal-to-noise diversity of the training gradient. Canonical DDPM training
samples a distinct timestep per batch element.

* INoiseScheduler.AddNoiseBatched(cleanBatch, noiseBatch, timesteps) — default
  interface implementation delegates to the scalar AddNoise per element for
  backward compatibility. Concrete schedulers can override with a fused
  batched implementation.
* NoisePredictorBase.PredictNoiseBatched(noisyBatch, timesteps, conditioning) —
  virtual with default slice-then-call-scalar implementation. Subclasses that
  want a fused batched forward override this to keep the training loop on-device.
* DiffusionModelBase.PredictNoiseBatched — model-level counterpart with the
  same default slice-then-call-scalar behavior.
* DiffusionModelBase.Train detects batched vs rank-1 input and routes through
  the batched noise scheduler + batched predictor when input.Rank >= 2. Rank-1
  unbatched inputs stay on the scalar path for backward compatibility.

This is the foundation for the industry-exceeding "batched-per-element +
fused-resident" diffusion training path — see follow-up commits that wire
DiffusionModelBase.Train through the fused compiled plan.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): wGAN-GP correctness fix + fused disc + DP-SGD wire-ups (4F/4G/4H)

Consumes Tensors PR ooples/AiDotNet.Tensors#763 (compiled backward with
createGraph=true, DP-SGD helper, multi-slot persistent inputs).

## 4G — WGAN-GP correctness fix + fused disc steps (5 files)

Root cause fixed for AiDotNet issue #1844: every ComputeGradientPenalty
implementation used the inner GradientTape's ComputeGradients WITHOUT
createGraph=true. The inner backward's ops didn't record on the outer
tape, so inputGradients had no GradFn chain back to the discriminator
weights. Effect: the gradient penalty was included in the loss VALUE
but NOT in the disc's gradient — WGAN-GP silently degraded to plain
WGAN with no 1-Lipschitz enforcement.

Fix (5 files: CTGAN, DPCTGAN, CopulaGAN, CausalGAN, TableGAN): add
createGraph: true to the inner ComputeGradients so the outer tape can
differentiate the penalty through the disc weights.

Also wired 4 disc steps through the fused compiled plan
(CausalGAN, CTGAN, CopulaGAN, TableGAN — DPCTGAN's disc is DP so goes
via 4H). Packs (real, fake) into one persistent input along axis 0;
loss closure splits scores back out and computes wasserstein + λ·GP.
The fused path activates once Tensors PR #763 lands (its 4C change
removes the !createGraph gate at GradientTape.cs:727); until then
these wire-ups fall through to the (now-corrected) eager tape path.

## 4H — DP-SGD wire-ups (2 files)

Local AiDotNet.Training.DpSgdStep<T> — drop-in mirror of the same
helper in Tensors PR #763 (Engines.Training.DpSgdStep<T>). Enables
the wire-ups to land in this PR without waiting on the Tensors NuGet
publish. Both implementations enforce the Abadi 2016 §3 Algorithm 1
clip-BEFORE-aggregate order via their structure.

* DPCTGAN.TrainDiscriminatorStepBatchedDP — routes the per-example
  WGAN-GP + DP-SGD critic through DpSgdStep<T>. Objective: Wasserstein
  + λ·GP per example, clipped per-example, aggregated, noised, averaged.
* MedSynth.TrainDiscriminatorStepPerExampleDPSGD — routes the per-
  example non-saturating BCE + DP-SGD critic through the same helper.

Once Tensors PR #763 merges + NuGet publishes, the local mirror can
be swapped for the Tensors version by changing the `using` — the API
surface is identical by design.

## 4F — Batched-per-element diffusion foundation

DiffusionModelBase.Train samples per-batch-element timesteps for
rank ≥ 2 inputs (Ho et al. 2020 canonical pattern; HuggingFace
diffusers reference), routing through NoiseSchedulerBase.AddNoiseBatched
+ DiffusionModelBase.PredictNoiseBatched. Rank-1 unbatched inputs
keep the scalar path for backward compat.

net471 doesn't support default interface implementations, so
AddNoiseBatched lives on NoiseSchedulerBase<T> (as virtual with a
default per-element delegate) rather than on the INoiseScheduler<T>
interface. DiffusionModelBase's fused-training wire-up (via Tensors
PR #763's PersistentInputRegistry) is queued for the NuGet-bump
follow-up commit — the API foundation is already in place.

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* fix(review): resolve 12 CodeRabbit findings on PR #1843

Addresses every unresolved review comment on the non-TS GPU-residency PR.

## Correctness fixes

* DiffusionModelBase.Train — RecomputeForward closure now dispatches
  isBatched ? PredictNoiseBatched(inp, timesteps) : PredictNoise(inp, timestep),
  so optimizer re-evaluation matches the recorded batched forward objective.
* NoisePredictorBase.PredictNoiseBatched — detect batch-aligned conditioning
  (leading dim == batchSize) and slice per element; preserve shared conditioning
  when the leading dim doesn't match. Fixes shape-mismatch / conditioning-shared
  bugs in classifier-free-guidance-style predictors.
* CSDI.Train — remove the fused-plan branch. ComputeDenoisingPairTape samples
  fresh (timestep, noise) each call, and calling it in both Fwd and Loss produces
  independent samples that don't match. The compiled plan can't refresh the RNG
  per replay either. Path stays on the eager tape until Tensors' PersistentInputRegistry
  (PR ooples/AiDotNet.Tensors#763) lands.
* TFC.Train — ComputeContrastiveLossTape now runs INSIDE ForwardCombined so it
  consumes the current-step persistent input (`inp`), not the outer `input` which
  would freeze at compile time. Closure-captured local; Loss reads it with a
  null-guard covering the Fwd-then-Loss ordering invariant.
* TOTEM.Train — capture the commitment tensor from ForwardNativeForTrainingWithCommitment's
  first call; reuse it in Loss instead of running the quantizer a second time.
  Without this, the compiled path performs an EMA SetCodebookValue update TWICE
  per step and diverges from the eager path.
* NoiseSchedulerBase.AddNoiseBatched — validate full shape parity (rank + every
  dim), not only the leading batch dim. Prevents indexing beyond noiseBatch's span
  when a caller passes [B, smaller...] shapes.

## Correctness cleanup

* CTGAN.TrainGeneratorStepBatched — capture (act, condFromInput, maskFromInput)
  from Fwd's single generator pass; reuse in Loss so the conditional-CE term
  doesn't re-run GeneratorForwardWithResidualBatched + ApplyOutputActivationsBatched
  (was doubling the per-step generator forward cost).
* TabSynGenerator.TrainVAEBatch — capture (mean, logVar) from Fwd's single encoder
  pass; reuse in Loss so the encoder doesn't run twice per row. Also: if the fused
  step fails mid-batch after fusedEngaged=true, drop to the eager path for the
  remaining rows (no row is silently skipped).
* FourierNeuralOperator.TapeTrainStep — remove the duplicate _fourierLayers loop
  in both the fused-path allTrainable collection AND the eager paramList. Fourier
  layers are already registered into Layers at construction, so iterating them
  again double-registered each Fourier parameter and drove the eager SGD to apply
  the update twice per step.

## API cleanup

* GpuResidentFusedStep<T> — public → internal. Aligns with CompiledTapeTrainingStep<T>
  which is already internal; keeps in-assembly training plumbing out of the public
  API surface.

Also fixed an inadvertent null-forgiving operator (`!`) usage that slipped into
the review-fix edits. All new nullable annotations use explicit null-guards with
descriptive InvalidOperationException on invariant violation (per CLAUDE.md's
"never use null-forgiving operators" rule).

INoiseScheduler default-interface issue (net471 incompatibility) was already
fixed in an earlier commit (AddNoiseBatched moved to NoiseSchedulerBase).

Builds green net471 / net8.0 / net10.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(training): vectorized fused DP-SGD/WGAN-GP/multi-slot primitives + consumer wire-ups

Adds three fused-plan training primitives (mirrors of AiDotNet.Tensors PR #763)
and wires the DP-SGD consumers through them:

* DpSgdFusedStep<T> — Abadi 2016 §3 Algorithm 1 per-example clip-before-aggregate.
  Runs each per-example forward+backward through a compiled plan (LR=0 so weights
  don't drift between replays), then computes the global L2 norm, clips, noises,
  and aggregates via vectorized IEngine ops (TensorMultiply/ReduceSum for the
  norm, TensorMultiplyScalar/TensorAdd for the accumulator, TensorRandomNormalInto
  for on-device Gaussian noise). Returns aggregated gradients so the caller's
  configured optimizer (Adam/AdamW/SGD/...) applies the update — no hardcoded
  SGD step inside.

* WganGpFusedStep<T> — Gulrajani 2017 WGAN-GP critic step. Composes
  E[D(fake)] − E[D(real)] + λ·(‖∇_x̃ D(x̃)‖₂ − 1)² inside one compiled plan with
  the inner ∇_x̃ D(x̃) recorded via createGraph=true so it differentiates into
  disc weights (issue #1844 fix). OnesLike uses vectorized Engine.TensorFill.

* MultiSlotFusedStep<T> — N-slot persistent input mechanism with plan cache
  keyed by composite shape + parameter identity. Slots refreshed via
  AsSpan().CopyTo(AsWritableSpan()) — no per-element loops on the hot path.

Consumer wire-ups (this PR):

* DPCTGANGenerator.TrainDiscriminatorStepBatchedDP — primary path routes through
  DpSgdFusedStep.TryStep, falls back to the existing ComputePerExampleNoisedGradients
  when the fused path can't engage (non-GPU host / compilation disabled).
* MedSynthGenerator.TrainDiscriminatorStepPerExampleDPSGD — same primary/fallback
  layering.
* Both legacy fallback loops (ComputePerExampleNoisedGradients + MedSynth eager)
  are now themselves vectorized: global-L2-norm via ReduceSum(g·g), clipped
  accumulation via TensorMultiplyScalar+TensorAdd, on-device Gaussian noise via
  TensorRandomNormalInto+TensorAdd — no scalar per-element loops anywhere on the
  DP-SGD path.

Codebase convention adopted:

* Class-scope `private static IEngine Engine => AiDotNetEngine.Current;` and
  `private static readonly INumericOperations<T> Ops = MathHelper.GetNumericOperations<T>();`
  on each fused-step class (matches the ~20 activation/optimizer/etc. bases).
  No IEngine or INumericOperations<T> threaded through method signatures.

Skipped (follow-ups filed):

* WGAN-GP consumer rewire → #1845 (blocked on IGradientBasedOptimizer<T> config
  accessor extension; consumers already correct via GpuResidentFusedStep + local
  createGraph=true ComputeGradientPenalty).
* Diffusion consumer rewire → #1846 (blocked on per-generator forward-refactor
  to move _timestepProjection inside compiled plan while treating raw timestep
  as a persistent slot).

Verification:

* net8.0 + net471 + net10.0 all build clean on both AiDotNet and AiDotNet.Tensors.
* Old DpSgdStep.cs deleted (replaced by DpSgdFusedStep).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* fix(review): tFC + TOTEM preprocessing rewritten with traceable engine ops

Resolves the two BLOCKING CodeRabbit findings on PR #1843: both fused-plan
training paths were freezing preprocessing into the compiled plan because
their inner ops used host-side .Data.Span loops, so later replays reused the
first batch's normalized values / spectrum / argmin decisions instead of
recomputing per batch.

TFC — ApplyInstanceNormalization (RevIN) + ComputeFrequencyRepresentation:

* ApplyInstanceNormalization now delegates to a new stateless NormalizeWithStats
  that returns (normalized, mean, std) as tensors. Under the hood: ReduceMean +
  ReduceVariance + TensorSqrt + TensorBroadcastSubtract + TensorBroadcastDivide.
  All ops record on the tape and re-execute per replay under the compiled plan.
* DenormalizeForecast likewise delegates to DenormalizeForecastWithStats, which
  takes mean/std as explicit tensor parameters. ForwardNative threads them
  through as locals so the compile-mode replay uses CURRENT-step stats instead
  of frozen trace-time values.
* _revinMean/_revinStd (Vector<T> scalars) → _revinMeanTensor/_revinStdTensor
  (Tensor<T>? nullable) — kept for the abstract override's external callers,
  but the fused path never reads them.
* ComputeFrequencyRepresentation now uses Engine.RFFT for batched real FFT →
  reshape to [B, halfN, 2] pairs → TensorMultiply + ReduceSum(axis=2) for
  magSquared → TensorSqrt + TensorMultiplyScalar(1/n) for the one-sided
  magnitude spectrum → TensorSlice + TensorFlip + TensorConcatenate to mirror
  bins [1..n-halfN] into the tail. Handles even and odd n identically to the
  old scalar impl.
* Fused-plan fast path in TFC.Train restored — now safe because both
  preprocessing methods trace correctly.

TOTEM — VectorQuantize (VQ-VAE) traceable rewrite + post-Step EMA:

* New VectorQuantizeTraceable returns (quantized, commitmentLoss, argmin, head)
  using Engine ops end-to-end: TensorBroadcastSubtract + TensorMultiply +
  ReduceSum for distances, TensorArgMin along the codebookSize axis, per-c
  TensorSliceAxis + TensorIndexSelectDiff + TensorStack for the gather,
  Engine.StopGradient for the straight-through estimator, ReduceSum-based
  commitment loss weighted by β/totalLen.
* EMA moved OUT of the compiled forward into a new UpdateCodebookEMA(head,
  argmin). Called POST-Step by the fused path with the trace-time graph-node
  references — their .Data reflects the LAST replay so the update lands
  exactly once per batch (matches CodeRabbit's "EMA must execute exactly once
  per batch" contract). Under the compiled plan, argmin/head are refreshed
  by every _plan.Step() so post-Step reads see the current batch's values.
* UpdateCodebookEMA expresses the per-codebook scatter as: current codebook
  slice + TensorScatterAdd((1-decay)·(head - gathered), argmin) → new slice,
  then TensorConcatenate across codebook axis into the full [numCodebooks,
  codebookSize, codebookDim] tensor, then Engine.TensorCopy back into the
  _codebooks tensor object to preserve identity (future reads via the same
  reference see the update).
* Legacy VectorQuantize is now a thin adapter around VectorQuantizeTraceable
  for callers that don't need the extras.
* ForwardNativeForTrainingWithCommitment delegates to a new
  ForwardNativeForTrainingWithVQExtras that exposes the argmin/head; the
  original (forecast, commitmentLoss) contract is preserved for callers that
  don't need EMA state.
* Fused-plan fast path in TOTEM.Train restored — now safe because VQ is
  fully traceable AND the EMA runs exactly once per batch in post-Step eager
  code.

No new Tensors primitives were needed — the engine already has RFFT,
ReduceMean/Variance, TensorSqrt, TensorBroadcastSubtract/Divide, TensorFlip,
TensorConcatenate, TensorArgMin, TensorIndexSelectDiff, TensorSliceAxis,
TensorStack, StopGradient, TensorScatterAdd, and TensorCopy, all in the
autodiff/compile registry per OpRegistry.cs.

Verification:

* net8.0 + net471 + net10.0 all build clean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* 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>

* feat(training): wire additional WGAN-GP + diffusion consumers through fused primitives (#1848)

Extends PR #1847 with three additional consumer wire-ups discovered during a
comprehensive audit against the primitives added in PR #1843. Each consumer
attempts the fused path first via the existing IFusedOptimizerSpec /
TryMapToFusedOptimizerConfig plumbing and falls back to eager tape training
when the fused path can't engage.

## WGAN-GP consumers (additional beyond #1845)

* NeuralNetworks/WGANGP.cs — standalone WGAN-GP class. Previously used the
  legacy flat-vector round-trip (three Critic.Predict calls per step,
  GetParameterGradients + host-side vector combine + UpdateCriticWithOptimizer).
  TrainCriticBatchWithGP now attempts WganGpFusedStep.TryStep first, using
  Critic.Layers as the parameter source and a per-batch ε ~ U(0, 1) sampler
  matching each critic's local ComputeGradientPenalty behavior. Falls back to
  the legacy path when the critic's optimizer has no fused-kernel mapping.

## Diffusion consumers (additional beyond #1846)

* NeuralNetworks/SyntheticData/FinDiffGenerator.cs — per-row DDPM training
  (Sattarov et al. 2023). TrainBatch caches a MultiSlotFusedStep across rows
  so the compiled plan is built once and replayed via slot-data refresh for
  subsequent rows. Slots: (packedDenoiserInput, targetNoise). Falls back to
  the existing Train(input, targetNoise) call on miss.

* NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs — per-row DDPM training
  with a custom TapeStepOver optimizer step. Same MultiSlotFusedStep caching
  pattern as FinDiff. Slots: (denoiserInput, targetNoise). Falls back to the
  existing tape-based TapeStepOver path on miss.

## Explicitly not wired in this PR (documented)

* NeuralNetworks/GenerativeAdversarialNetwork.cs — uses BCE loss with optional
  GP regularization (a separate auxiliary optimizer step), NOT Wasserstein +
  GP. Wiring WganGpFusedStep would change training semantics from BCE to
  Wasserstein. Requires a separate design decision.

* Finance/Forecasting/Foundation/{CCDM,MGTSD,TSDiff,TimeDiff,TimeGrad}.cs and
  Finance/Probabilistic/DiffusionTS.cs — each uses .Data.Span host-side packs
  for the denoising input (same trace-freeze issue TFC/CSDI hit in PR #1843).
  Each needs a TFC/CSDI-scale traceable rewrite BEFORE fused wiring is safe.
  Deferred to individual per-consumer PRs.

* MetaLearning/Algorithms/{MetaDDPMAlgorithm,MetaDMAlgorithm,MetaDiffAlgorithm}.cs
  — hand-rolled Vector&lt;T&gt; params with index arithmetic, NOT the
  ILayer/Engine/tape infrastructure. Would need full rewrite to fit MultiSlotFusedStep.

* DiffusionModelBase subclasses (DDPMModel, DiffWaveModel, LatentDiffusionModelBase)
  — the SupportsFusedDenoising opt-in hook (added in PR #1847) is available,
  but flipping it safely requires per-class audit of PredictNoise → UNet
  traceability. Left as follow-up work; the mechanism is in place.

## Verification

* net8.0, net471, net10.0 all build clean.
* No API surface changes.

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

* chore(deps): bump AiDotNet.Tensors 0.111.2 -> 0.113.0 (activates #763 for fused paths)

0.113.0 publishes Tensors PR #763 — this branch's gating dependency. #763's key fix
(4C): the compiled/persistent backward now honors createGraph=true (GradientTape.
ComputeGradients previously gated the compiled path on !createGraph), so the WGAN-GP
gradient penalty's inner backward differentiates into the disc weights through the
fused compiled plan instead of silently returning zeros (issue #1844). Against 0.111.2
the fused GPU-resident WGAN-GP path would have silently degraded to plain WGAN.

The src/Training fused-primitive mirrors (WganGpFusedStep, MultiSlotFusedStep,
DpSgdFusedStep) stay as the consumer-side primitives that #1847/#1848 centralize every
consumer through; they call the public engine API, so the bump alone routes them onto
#763's fixed compiled backward. Also picks up #765 (compiled ReduceMax axis fill) and
#764 (resident-param fused-Adam fix). Native OneDNN/OpenBLAS/CLBlast bumped in lockstep
(all published). Builds green on net8.0.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

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

Author: 3 people

Directory.Packages.props

@@ -235,11 +235,24 @@
          accumulating multi-consumer grad buffers) and #764 (GPU-resident-parameter fused-Adam mistrain
          fix), which together make this PR's TimeSeries GPU-resident training path train faithfully (the
          resident run now lowers the eager-forward training MSE and improves held-out MSE, so the N-BEATS
-         correctness gate accepts it). 0.112.0 is PUBLISHED to NuGet and is a superset of 0.111.2. -->
-    <PackageVersion Include="AiDotNet.Tensors" Version="0.112.0" />
-    <PackageVersion Include="AiDotNet.Native.OneDNN" Version="0.112.0" />
-    <PackageVersion Include="AiDotNet.Native.OpenBLAS" Version="0.112.0" />
-    <PackageVersion Include="AiDotNet.Native.CLBlast" Version="0.112.0" />
+         correctness gate accepts it). 0.112.0 is PUBLISHED to NuGet and is a superset of 0.111.2.
+
+         Bumped 0.112.0 -> 0.113.0: brings AiDotNet.Tensors #763 — this branch's GATING dependency and
+         the engine half of the non-time-series GPU-residency sweep. #763's key fix (4C): the
+         compiled/persistent backward now honors createGraph=true (GradientTape.ComputeGradients
+         previously gated the compiled path on !createGraph), so the WGAN-GP gradient penalty's inner
+         backward differentiates into the disc weights through the fused compiled plan instead of
+         silently returning zeros (issue #1844). Against 0.112.0 the fused GPU-resident WGAN-GP path
+         (WganGpFusedStep / GpuResidentFusedStep) would have silently degraded to plain WGAN. #763 also
+         publishes Engines.Training.{DpSgdStep, DpSgdFusedStep, MultiSlotFusedStep, WganGpFusedStep,
+         PersistentInputRegistry}; the src/Training mirror classes stay as the consumer-side primitives
+         (the fused-primitive centralization in #1847/#1848 wires every consumer through them) and now
+         run against the fixed engine. 0.113.0 is PUBLISHED to NuGet and is a superset of 0.112.0 (also
+         carries #765 compiled-ReduceMax axis fill), so all rationale above is retained. -->
+    <PackageVersion Include="AiDotNet.Tensors" Version="0.113.0" />
+    <PackageVersion Include="AiDotNet.Native.OneDNN" Version="0.113.0" />
+    <PackageVersion Include="AiDotNet.Native.OpenBLAS" Version="0.113.0" />
+    <PackageVersion Include="AiDotNet.Native.CLBlast" Version="0.113.0" />
     <!-- Microsoft / .NET -->
     <PackageVersion Include="Microsoft.CSharp" Version="4.7.0" />
     <PackageVersion Include="Microsoft.Data.Sqlite" Version="10.0.9" />

src/Audio/Fingerprinting/CLAPModel.cs

@@ -501,6 +501,52 @@ public override void Train(Tensor<T> input, Tensor<T> expected)
 
             var optimizer = GetOrCreateBaseOptimizer();
 
+            // GPU-RESIDENT fast path — audio + text encoders + the learned
+            // temperature scalar. _logTemperature is NOT an ITrainableLayer, so
+            // it goes through extraTensors (Phase 4A). The fused optimizer tracks
+            // it via the same moment-buffer path as layer-carried params.
+            var trainableLayers = new List<ITrainableLayer<T>>();
+            foreach (var l in Layers) if (l is ITrainableLayer<T> t) trainableLayers.Add(t);
+            foreach (var l in TextEncoderLayers) if (l is ITrainableLayer<T> t) trainableLayers.Add(t);
+            var extras = new List<Tensor<T>> { _logTemperature };
+            if (trainableLayers.Count > 0 || extras.Count > 0)
+            {
+                // The contrastive loss depends on BOTH encoders + temperature, so
+                // the forward closure runs the full symmetric-alignment computation
+                // and the loss closure just reduces its output. We pass the audio
+                // input as the primary and the text batch as target so the closure
+                // signature matches TryStep's contract.
+                Tensor<T> FwdCLAP(Tensor<T> audioIn) => EncodeAudio(audioIn); // audio embedding — the loss closure re-runs the text side.
+                Tensor<T> LossCLAP(Tensor<T> audioEmb, Tensor<T> textBatch)
+                {
+                    // Re-run text encoder on this replay's text batch.
+                    var textEmb = EncodeText(textBatch);
+                    int batchSize = audioEmb.Shape[0];
+                    int projDim = audioEmb.Shape[audioEmb.Shape.Length - 1];
+                    var audioEmb2D = audioEmb.Shape.Length == 2 ? audioEmb : Engine.Reshape(audioEmb, new[] { batchSize, projDim });
+                    var textEmb2D = textEmb.Shape.Length == 2 ? textEmb : Engine.Reshape(textEmb, new[] { batchSize, projDim });
+                    var textEmbT = Engine.TensorTranspose<T>(textEmb2D);
+                    var sim = Engine.TensorMatMul<T>(audioEmb2D, textEmbT);
+                    var tau = Engine.TensorExp<T>(_logTemperature);
+                    var tauBroadcast = Engine.TensorTile(Engine.Reshape(tau, new[] { 1, 1 }), new[] { batchSize, batchSize });
+                    var logitsA2T = Engine.TensorMultiply<T>(sim, tauBroadcast);
+                    var logitsT2A = Engine.TensorTranspose<T>(logitsA2T);
+                    var halfA2T = SymmetricRowCrossEntropy(logitsA2T, batchSize);
+                    var halfT2A = SymmetricRowCrossEntropy(logitsT2A, batchSize);
+                    return Engine.TensorAdd<T>(halfA2T, halfT2A);
+                }
+                if (AiDotNet.Training.GpuResidentFusedStep<T>.TryStep(
+                        trainableLayers, input, expected,
+                        forward: FwdCLAP, computeLoss: LossCLAP,
+                        optimizer: optimizer,
+                        out T fusedLoss,
+                        extraTensors: extras))
+                {
+                    LastLoss = fusedLoss;
+                    return;
+                }
+            }
+
             using var tape = new GradientTape<T>();
             // Forward both encoders inside the same tape so gradients flow
             // through every parameter. EncodeAudio / EncodeText already

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>
@@ -802,6 +825,41 @@ protected virtual Tensor<T> Generate(int[] shape, int numInferenceSteps, int? se
     /// <inheritdoc />
     public abstract Tensor<T> PredictNoise(Tensor<T> noisySample, int timestep);
 
+    /// <summary>
+    /// Batched per-element noise prediction (industry-standard DDPM training pattern).
+    /// Default implementation slices the batch and calls scalar <see cref="PredictNoise"/>
+    /// per element — subclasses should override with a fused batched forward to keep
+    /// training on-device. <paramref name="noisyBatch"/> is shape <c>[B, ...]</c>;
+    /// <paramref name="timesteps"/> is a <c>[B]</c> int vector.
+    /// </summary>
+    public virtual Tensor<T> PredictNoiseBatched(Tensor<T> noisyBatch, int[] timesteps)
+    {
+        int batchSize = noisyBatch.Shape[0];
+        if (timesteps.Length != batchSize)
+            throw new ArgumentException(
+                $"timesteps length {timesteps.Length} does not match batch size {batchSize}.",
+                nameof(timesteps));
+
+        int perElement = noisyBatch.Length / batchSize;
+        var elemShape = new int[noisyBatch.Rank - 1];
+        for (int i = 1; i < noisyBatch.Rank; i++) elemShape[i - 1] = noisyBatch.Shape[i];
+        var result = new Tensor<T>(noisyBatch._shape);
+        var nbSpan = noisyBatch.AsSpan();
+        var resSpan = result.AsWritableSpan();
+        for (int b = 0; b < batchSize; b++)
+        {
+            var elem = new Tensor<T>(elemShape);
+            var elemSpan = elem.AsWritableSpan();
+            for (int j = 0; j < perElement; j++)
+                elemSpan[j] = nbSpan[b * perElement + j];
+            var pred = PredictNoise(elem, timesteps[b]);
+            var predSpan = pred.AsSpan();
+            for (int j = 0; j < perElement; j++)
+                resSpan[b * perElement + j] = predSpan[j];
+        }
+        return result;
+    }
+
     /// <summary>
     /// Runs one denoising-step noise prediction, optionally inside a GPU deferred execution graph
     /// (AiDotNet.Tensors #642) when <see cref="DiffusionModelOptions{T}.UseGpuExecutionGraph"/> is
@@ -1034,19 +1092,143 @@ public virtual void Train(Tensor<T> input, Tensor<T> expectedOutput)
         // tensors than GetParameters knows about, which is now the norm after
         // migrating layers like FlashAttentionLayer from Matrix<T> to Tensor<T>.
 
-        // Sample a random timestep and build the noisy training sample.
-        var timestep = RandomGenerator.Next(_scheduler.Config.TrainTimesteps);
-        var inputVector = input.ToVector();
-        var noiseVector = SampleNoise(inputVector.Length, RandomGenerator);
-        var noisySample = _scheduler.AddNoise(inputVector, noiseVector, timestep);
-        var noisySampleTensor = new Tensor<T>(input._shape, noisySample);
+        // Industry-standard batched-per-element timesteps (Ho et al. 2020, HuggingFace
+        // diffusers reference): sample a distinct timestep per batch element instead of
+        // one timestep for the whole batch. Decorrelates the noise-schedule signal
+        // across the batch, which is the canonical DDPM training pattern.
+        //
+        // Rank-1 input (unbatched, historical AiDotNet contract) still gets a single
+        // timestep; rank ≥ 2 gets per-element timesteps.
+        bool isBatched = input.Rank >= 2;
+        int batchSize = isBatched ? input.Shape[0] : 1;
+        var timesteps = new int[batchSize];
+        for (int b = 0; b < batchSize; b++)
+            timesteps[b] = RandomGenerator.Next(_scheduler.Config.TrainTimesteps);
+        // Legacy scalar view — retained for downstream code that reads the "current"
+        // timestep (e.g. QAT hook telemetry). For batched inputs this reports element 0's
+        // timestep, matching the historical single-timestep-per-Train contract.
+        var timestep = timesteps[0];
+
+        Tensor<T> noisySampleTensor;
+        Vector<T> noiseVector;
+        if (isBatched)
+        {
+            var noiseBatch = new Tensor<T>(input._shape);
+            var noiseSpan = noiseBatch.AsWritableSpan();
+            for (int i = 0; i < noiseSpan.Length; i++)
+                noiseSpan[i] = NumOps.FromDouble(RandomGenerator.NextGaussian());
+            // AddNoiseBatched lives on NoiseSchedulerBase (not INoiseScheduler — that
+            // interface can't carry a default implementation on net471). Fall back to
+            // per-element scalar AddNoise if a caller passed a scheduler that doesn't
+            // derive from NoiseSchedulerBase (shouldn't happen for framework schedulers).
+            if (_scheduler is Schedulers.NoiseSchedulerBase<T> baseScheduler)
+            {
+                noisySampleTensor = baseScheduler.AddNoiseBatched(input, noiseBatch, timesteps);
+            }
+            else
+            {
+                noisySampleTensor = new Tensor<T>(input._shape);
+                var cleanSpan = input.AsSpan();
+                var noisedSpan = noisySampleTensor.AsWritableSpan();
+                int perElement = input.Length / batchSize;
+                for (int b = 0; b < batchSize; b++)
+                {
+                    var cleanVec = new Vector<T>(perElement);
+                    var noiseVec = new Vector<T>(perElement);
+                    for (int j = 0; j < perElement; j++)
+                    {
+                        cleanVec[j] = cleanSpan[b * perElement + j];
+                        noiseVec[j] = noiseSpan[b * perElement + j];
+                    }
+                    var noised = _scheduler.AddNoise(cleanVec, noiseVec, timesteps[b]);
+                    for (int j = 0; j < perElement; j++)
+                        noisedSpan[b * perElement + j] = noised[j];
+                }
+            }
+            noiseVector = noiseBatch.ToVector();
+        }
+        else
+        {
+            var inputVector = input.ToVector();
+            noiseVector = SampleNoise(inputVector.Length, RandomGenerator);
+            var noisySample = _scheduler.AddNoise(inputVector, noiseVector, timestep);
+            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
         // trainable parameters. Collection must happen AFTER the forward pass so
-        // newly-initialized layers are visible to the walker.
-        var predicted = PredictNoise(noisySampleTensor, timestep);
+        // newly-initialized layers are visible to the walker. Batched inputs route
+        // through the per-element PredictNoiseBatched (Ho et al. 2020 canonical pattern);
+        // rank-1 unbatched inputs go through the scalar PredictNoise for backward compat.
+        var predicted = isBatched
+            ? PredictNoiseBatched(noisySampleTensor, timesteps)
+            : PredictNoise(noisySampleTensor, timestep);
         var paramTensors = CollectTrainableParameters();
         if (paramTensors.Length == 0)
         {
@@ -1150,7 +1332,10 @@ public virtual void Train(Tensor<T> input, Tensor<T> expectedOutput)
         // clip; nothing is materialized or copied here. The forward/loss closures are only consulted by
         // optimizers that re-evaluate the objective (e.g. line search); Adam ignores them.
         T lossValue = loss.Length > 0 ? loss[0] : NumOps.Zero;
-        Tensor<T> RecomputeForward(Tensor<T> inp, Tensor<T> _) => PredictNoise(inp, timestep);
+        Tensor<T> RecomputeForward(Tensor<T> inp, Tensor<T> _) =>
+            isBatched
+                ? PredictNoiseBatched(inp, timesteps)
+                : PredictNoise(inp, timestep);
         Tensor<T> RecomputeLoss(Tensor<T> inp, Tensor<T> target)
         {
             using var noGrad = new NoGradScope<T>();