fix: 140x optimizer speedup — lazy stats, in-place updates, skip redundant Train()

[ooples]

Summary

Fixes #1123 — critical performance bug where optimizer training took 5+ minutes instead of <1 second.

Before: 500 samples × 100 iterations = >30s timeout
After: 214ms (140x speedup)

Root Causes Found & Fixed

1. Cache key used object hash (never hit)

GenerateCacheKey used parameters.GetHashCode() (memory address). Since UpdateSolution creates new objects each epoch, the cache key changed every time → model.Train() ran every epoch redundantly.

Fix: Content-based hash sampling actual parameter values.

2. WithParameters did Serialize+Deserialize 1600x

UpdateSolution called WithParameters → Clone → DeepCopy → Serialize+Deserialize for every batch update (16 batches × 100 epochs = 1600 full serialization roundtrips).

Fix: In-place SetParameters for gradient optimizers. Lightweight CreateNewInstance instead of Clone for WithParameters.

3. Full stats computed every epoch (30+ metrics × 3 datasets)

EvaluateModelDirectly eagerly created ErrorStats (20 metrics), PredictionStats (30 metrics), BasicStats (13 metrics), ModelStats (correlation matrix, VIF, LOO, posterior sampling — O(n²×features)) for ALL 3 datasets, every epoch. Only R² was needed.

Fix: All 4 stats classes are now fully lazy — properties compute on first access. Only PredictionStats.R2 is computed eagerly. EvaluateModelDirectly only computes the dataset the FitnessCalculator.PreferredDataSetType needs.

4. SelectFeatures allocated matrix even for all features

When all features selected (the common case), SelectFeaturesMatrix allocated a full copy. Now returns the original.

5. Gradient optimizers skip Train() after initial training

GradientBasedOptimizerBase sets SkipTrainingInEvaluation=true after first evaluation, since parameter updates happen via UpdateSolution, not Train().

Files Changed (11)

• src/Optimizers/OptimizerBase.cs — cache key, skip training, lazy eval
• src/Optimizers/GradientBasedOptimizerBase.cs — in-place SetParameters, skip flag
• src/Statistics/PredictionStats.cs — fully lazy (R² eager, rest deferred)
• src/Statistics/ErrorStats.cs — fully lazy
• src/Statistics/BasicStats.cs — fully lazy
• src/Statistics/ModelStats.cs — fully lazy
• src/Regression/RegressionBase.cs — lightweight WithParameters
• src/Helpers/OptimizerHelper.cs — SelectFeatures short-circuit
• src/Interfaces/IFitnessCalculator.cs — PreferredDataSetType
• src/FitnessCalculators/FitnessCalculatorBase.cs — expose DataSetType
• tests/.../OptimizerTrainSkipTests.cs — 3 integration tests

Test plan

• 500 samples × 100 iterations completes in 214ms (was >30s)
• Gradient optimizer calls Train() at most 2x (not 1000x)
• Cache prevents redundant evaluation
• Existing optimizer tests pass
• ModelFamily regression tests pass

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added tensor-level overrides for 40+ activation functions enabling gradient-tape recording for autodiff
  • Introduced CrossEntropyWithLogitsLoss for raw logit-based classification
  • Added PReLULayer and ParallelStreamsLayer for extended neural architecture support
  • Enhanced test diagnostics with hung-test detection

• Improvements

  • Optimized tensor shape handling by eliminating unnecessary array copies
  • Migrated reshape/transpose operations to engine-based implementations for better autodiff integration
  • Refactored vectorized operations reducing scalar loops in training pipelines
  • Improved model persistence guards and parameter handling workflows

• Chores

  • Updated AiDotNet.Tensors dependency to 0.37.0
  • Expanded code coverage exclusions for internal generators and tooling

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of files (300). Try reducing the number of changed files and requesting a review from Copilot again.

[coderabbitai]

Actionable comments posted: 20

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (7)

src/NeuralNetworks/Layers/MaskingLayer.cs (1)

378-384: 🧹 Nitpick | 🔵 Trivial

Consider unifying dispose patterns for consistency.

The ResetState method uses a defensive is IDisposable check, while the newly added disposal in ForwardGpu (line 238) uses the simpler ?.Dispose() pattern. Since Tensor<T> appears to have a public Dispose() method (evidenced by line 238 working), you could simplify ResetState to match:

♻️ Optional consistency improvement

         // Clear GPU mask tensor
-        if (_lastMaskGpu is IDisposable disposable)
-        {
-            disposable.Dispose();
-        }
+        _lastMaskGpu?.Dispose();
         _lastMaskGpu = null;

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/MaskingLayer.cs` around lines 378 - 384, In
ResetState, simplify the disposal of the GPU mask to match the pattern used in
ForwardGpu: call _lastMaskGpu?.Dispose() and then set _lastMaskGpu = null;
instead of the current defensive "is IDisposable" check—update ResetState to use
the cleaner nullable Dispose pattern for _lastMaskGpu (same disposal semantics
as ForwardGpu).

src/NeuralNetworks/Layers/GraphTransformerLayer.cs (1)

1341-1355: 🧹 Nitpick | 🔵 Trivial

Minor inconsistency in shape access, but functionally correct.

Line 1351 now uses bias._shape directly for iterating dimensions, while line 1349 still uses bias.Shape.Length. This works correctly but creates a subtle inconsistency — consider using bias._shape.Length on line 1349 for consistency, or documenting why the difference exists.

♻️ Optional: Consistent shape access

         if (hasBias)
         {
             var bias = _structuralBias ?? throw new InvalidOperationException("Structural bias is null during serialization.");
-            writer.Write(bias.Shape.Length);
+            writer.Write(bias._shape.Length);
             foreach (var dim in bias._shape) writer.Write(dim);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/GraphTransformerLayer.cs` around lines 1341 - 1355,
In GraphTransformerLayer.Serialize, the code inconsistently reads the tensor
shape via bias.Shape.Length but iterates bias._shape; make access consistent by
replacing bias.Shape.Length with bias._shape.Length (or vice versa) so both the
length check and the subsequent foreach use the same underlying representation
(refer to method Serialize and the local variable bias); ensure you keep the
existing null guard and exception behavior unchanged.

src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs (2)

627-635: ⚠️ Potential issue | 🟡 Minor

Inconsistent null handling for bias gradients vs weight gradients.

When _weightsGradient is null, the index advances by OutputFeatures * InputFeatures (line 624), leaving zeros in the gradient vector. However, when _biasesGradient is null, the index does NOT advance, meaning the returned vector is undersized or subsequent reads are off.

Suggested fix for consistent null handling

         // Biases gradients: [OutputFeatures] — scalar bias per output
         if (_biasesGradient != null)
         {
             for (int o = 0; o < OutputFeatures; o++)
-                    gradients[idx++] = _biasesGradient[o];
+                gradients[idx++] = _biasesGradient[o];
+        }
+        else
+        {
+            idx += OutputFeatures;
         }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs` around lines 627 - 635,
The method currently advances idx when _weightsGradient is null but not when
_biasesGradient is null, causing mis-sized/shifted gradients; update the branch
that handles _biasesGradient (the block that checks _biasesGradient != null and
fills gradients[idx++]) to mirror the _weightsGradient null handling by
advancing idx by OutputFeatures when _biasesGradient is null so the final
gradients array length and positions (for weights: OutputFeatures *
InputFeatures, for biases: OutputFeatures) remain consistent; adjust references
to _biasesGradient, _weightsGradient, gradients, idx, OutputFeatures and
InputFeatures accordingly.

---

354-356: ⚠️ Potential issue | 🟡 Minor

Documentation claims Riemannian gradient descent but implementation uses Euclidean.

The docstring at lines 354-356 states "Uses Riemannian gradient descent (exponential map of negative gradient)", but the actual implementation (lines 546-552) performs standard Euclidean updates: w = w - lr * grad. Riemannian SGD would require the exponential map to project the updated parameters back onto the manifold.

Either update the documentation to reflect the actual behavior, or implement true Riemannian updates using the defined ExpMapFromOrigin helper.

Documentation fix if Euclidean descent is intentional

     /// <summary>
     /// Updates the parameters of the layer using the calculated gradients.
     /// </summary>
-    /// <remarks>
-    /// Uses Riemannian gradient descent (exponential map of negative gradient).
-    /// </remarks>
+    /// <remarks>
+    /// Uses standard Euclidean gradient descent. Weights are stored in tangent space
+    /// at the origin, so Euclidean updates are appropriate.
+    /// </remarks>
     /// <param name="learningRate">The learning rate to use for the update.</param>

Also applies to: 539-555

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs` around lines 354 - 356,
The docstring in HyperbolicLinearLayer claims "Uses Riemannian gradient descent
(exponential map of negative gradient)" but the update code uses Euclidean
updates (w = w - lr * grad); either change the comment to reflect Euclidean SGD
or implement proper Riemannian updates: compute the tangent update vector as -lr
* grad (using the same shape as weights) and apply the manifold exponential map
via ExpMapFromOrigin to move the weight onto the hyperbolic manifold (i.e.,
replace the Euclidean assignment with a call that maps the updated tangent
vector back to the manifold), updating the same weight variable (w) used in the
current update loop; ensure you modify the update occurring in the
HyperbolicLinearLayer method that currently performs w = w - lr * grad (lines
~546-552) and use ExpMapFromOrigin with the proper arguments instead of the
simple subtraction if you choose the Riemannian approach.

src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs (1)

657-665: ⚠️ Potential issue | 🟠 Major

Higher-rank output restoration still drops trailing output axes.

This branch rebuilds the output shape as [..., output.Shape[1]]. That works for 2D expert outputs, but it breaks the higher-rank path you just added in CombineExpertOutputs: an output shaped [flatBatch, c, h, w] would be restored as [..., c]. Build the restored shape from the original leading input dims plus all non-batch output dims.

Proposed fix

- int outputFeatures = output.Shape[1];
- int[] newShape = new int[_originalInputShape.Length];
- for (int d = 0; d < _originalInputShape.Length - 1; d++)
-     newShape[d] = _originalInputShape[d];
- newShape[_originalInputShape.Length - 1] = outputFeatures;
+ int leadingRank = _originalInputShape.Length - 1;
+ int trailingRank = output.Shape.Length - 1;
+ int[] newShape = new int[leadingRank + trailingRank];
+ for (int d = 0; d < leadingRank; d++)
+     newShape[d] = _originalInputShape[d];
+ for (int d = 0; d < trailingRank; d++)
+     newShape[leadingRank + d] = output.Shape[d + 1];
  output = Engine.Reshape(output, newShape);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs` around lines 657 - 665,
The current restoration uses only output.Shape[1], which drops any trailing
spatial/channel axes produced by CombineExpertOutputs; instead, rebuild newShape
by copying the original leading (non-feature) dims from _originalInputShape and
then append all non-batch dims from output.Shape (i.e., copy output.Shape[1..]
rather than just output.Shape[1]) before calling Engine.Reshape; locate the
reshaping block that references _originalInputShape, output and Engine.Reshape
and replace the single-dimension assignment with a loop that copies each
output.Shape[k] into the corresponding tail positions of newShape.

src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs (1)

765-778: ⚠️ Potential issue | 🟠 Major

Make UpdateParameters preserve tensor identity too.

SetParameters is now in-place, but PrimaryCapsuleLayer<T>.UpdateParameters(T learningRate) still replaces _convWeights and _convBias on Line 673 and Line 676. That leaves the registered trainable-parameter references stale again after the first gradient step, so this fix only holds for callers that go through SetParameters.

Concrete direction

         var scaledWeightGrad = Engine.TensorMultiplyScalar(_convWeightsGradient, learningRate);
-        _convWeights = Engine.TensorSubtract(_convWeights, scaledWeightGrad);
+        var updatedWeights = Engine.TensorSubtract(_convWeights, scaledWeightGrad);
+        updatedWeights.Data.Span.CopyTo(_convWeights.Data.Span);

         var scaledBiasGrad = Engine.TensorMultiplyScalar(_convBiasGradient, learningRate);
-        _convBias = Engine.TensorSubtract(_convBias, scaledBiasGrad);
+        var updatedBias = Engine.TensorSubtract(_convBias, scaledBiasGrad);
+        updatedBias.Data.Span.CopyTo(_convBias.Data.Span);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs` around lines 765 - 778,
Update PrimaryCapsuleLayer<T>.UpdateParameters so it updates the existing tensor
data in-place instead of reassigning _convWeights and _convBias (which breaks
registered trainable references). Locate UpdateParameters and replace any code
that creates new tensors or assigns new instances to _convWeights/_convBias with
logic that writes the updated values into _convWeights.Data.Span and
_convBias.Data.Span (or elementwise-update via the tensor's buffer) using the
computed gradients and the learningRate, mirroring the in-place approach used by
SetParameters; this preserves tensor identity while applying the parameter
update.

src/Diffusion/Control/ControlNetModel.cs (1)

894-911: ⚠️ Potential issue | 🟡 Minor

Add input validation to match documented contract.

The documentation specifies that input should be [C, H, W] spatial tensor, but the Encode method doesn't validate this. Invalid input shapes would produce confusing errors deep in the convolution layers rather than a clear validation error.

🛡️ Proposed fix: Add input validation

 public List<Tensor<T>> Encode(Tensor<T> controlImage)
 {
+    if (controlImage.Shape.Length != 3)
+    {
+        throw new ArgumentException(
+            $"Control image must be [C, H, W] tensor with 3 dimensions, got {controlImage.Shape.Length} dimensions.",
+            nameof(controlImage));
+    }
+    if (controlImage.Shape[0] != _inputChannels)
+    {
+        throw new ArgumentException(
+            $"Control image channels ({controlImage.Shape[0]}) must match expected input channels ({_inputChannels}).",
+            nameof(controlImage));
+    }
+
     var features = new List<Tensor<T>>();
     var x = controlImage;

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Diffusion/Control/ControlNetModel.cs` around lines 894 - 911, The Encode
method lacks input validation for the documented [C,H,W] spatial tensor
contract; before using _downBlocks and _zeroConvs, check that controlImage is a
3D tensor (rank/shape length == 3) and each dimension > 0, and if not throw an
ArgumentException (or ArgumentNullException if controlImage is null) with a
clear message indicating expected [C,H,W] format; place this validation at the
start of Encode(Tensor<T> controlImage) so invalid shapes fail fast and provide
an actionable error rather than surfacing deeper convolution errors.

♻️ Duplicate comments (12)

src/Audio/Enhancement/SpikingFullSubNet.cs (1)

306-313: ⚠️ Potential issue | 🔴 Critical

Blocking: enforce exact shape equality before masking.

This still accepts same-length tensors with incompatible dimensions; fail fast on full shape mismatch to avoid silent audio corruption.

Proposed fix

-        if (stft.Length != mask.Length)
-            throw new ArgumentException(
-                $"Mask length ({mask.Length}) must match STFT length ({stft.Length}). " +
-                "A mismatch indicates a model output shape bug upstream.");
+        if (stft._shape.Length != mask._shape.Length)
+            throw new ArgumentException(
+                $"Mask shape rank ({mask._shape.Length}) must match STFT rank ({stft._shape.Length}).");
+        for (int d = 0; d < stft._shape.Length; d++)
+        {
+            if (stft._shape[d] != mask._shape[d])
+                throw new ArgumentException(
+                    $"Mask shape ({string.Join('x', mask._shape)}) must match STFT shape ({string.Join('x', stft._shape)}). " +
+                    "A mismatch indicates a model output shape bug upstream.");
+        }

As per coding guidelines, missing validation of external inputs is a blocking production-readiness issue.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Audio/Enhancement/SpikingFullSubNet.cs` around lines 306 - 313, The code
only checks stft.Length vs mask.Length and then multiplies elementwise, which
can hide incompatible dimensions; update the validation in SpikingFullSubNet
(the block creating result = new Tensor<T>(stft._shape) and looping over
stft[i]) to compare the full shapes of stft and mask (e.g., their
_shape/Dimensions or Shape property) and throw an ArgumentException if they
differ, before allocating result and calling NumOps.Multiply over elements;
ensure you reference the same Tensor<T> shape representation used in the class
to perform an exact shape equality check rather than just comparing Length.

src/Audio/Enhancement/CMGAN.cs (1)

333-340: ⚠️ Potential issue | 🔴 Critical

Blocking: validate full tensor shape match, not only flattened length.

Length equality alone still permits mismatched dimensions (same element count, wrong layout), which can silently corrupt masking output.

Proposed fix

-        if (stft.Length != mask.Length)
-            throw new ArgumentException(
-                $"Mask length ({mask.Length}) must match STFT length ({stft.Length}). " +
-                "A mismatch indicates a model output shape bug upstream.");
+        if (stft._shape.Length != mask._shape.Length)
+            throw new ArgumentException(
+                $"Mask shape rank ({mask._shape.Length}) must match STFT rank ({stft._shape.Length}).");
+        for (int d = 0; d < stft._shape.Length; d++)
+        {
+            if (stft._shape[d] != mask._shape[d])
+                throw new ArgumentException(
+                    $"Mask shape ({string.Join('x', mask._shape)}) must match STFT shape ({string.Join('x', stft._shape)}). " +
+                    "A mismatch indicates a model output shape bug upstream.");
+        }

As per coding guidelines, missing validation of external inputs is a blocking production-readiness issue.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Audio/Enhancement/CMGAN.cs` around lines 333 - 340, The current check
only compares flattened Length, which allows mismatched dimensions with equal
element counts; update the validation in CMGAN (where stft, mask, result are
used) to verify the full tensor shapes are identical (compare the shape
arrays/dimensions of stft and mask, e.g., stft._shape vs mask._shape or use an
existing Tensor shape-equality helper) and throw an ArgumentException with a
clear message if they differ; only after this full-shape check create result =
new Tensor<T>(stft._shape) and perform elementwise multiplication using
NumOps.Multiply to ensure layout alignment.

src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs (1)

397-399: ⚠️ Potential issue | 🟠 Major

Hoist output-projection bias reshape out of the timestep loop.

Engine.Reshape(_outputProjectionBias, ...) on Line 398 is loop-invariant but recomputed every timestep, adding avoidable overhead in the hot path.

Proposed fix

+        var outputProjectionBiasReshaped = Engine.Reshape(_outputProjectionBias, new[] { 1, _modelDimension });
+
         for (int t = 0; t < seqLen; t++)
         {
             ...
             // Output projection
             var y_t = Engine.TensorMatMul(h_t, _outputProjectionWeights);
-            var outBias = Engine.Reshape(_outputProjectionBias, new[] { 1, _modelDimension });
-            y_t = Engine.TensorBroadcastAdd(y_t, outBias);
+            y_t = Engine.TensorBroadcastAdd(y_t, outputProjectionBiasReshaped);
             ...
         }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs` around lines 397 - 399,
The Engine.Reshape(_outputProjectionBias, new[] { 1, _modelDimension }) call is
loop-invariant and should be hoisted out of the timestep loop: compute and store
the reshaped bias once (e.g., var outBias =
Engine.Reshape(_outputProjectionBias, new[] { 1, _modelDimension }) or
outBiasReshaped) before iterating timesteps inside the ExtendedLSTMLayer method,
then reuse that variable inside the loop when computing y_t
(Engine.TensorMatMul(h_t, _outputProjectionWeights) followed by
Engine.TensorBroadcastAdd(y_t, outBias)); this removes the repeated reshape
overhead while preserving the {1, _modelDimension} shape.

src/Document/PixelToSequence/MATCHA.cs (1)

626-636: ⚠️ Potential issue | 🟠 Major

Hot-loop allocations still create avoidable GC pressure.

sliceData, channelTensor, centered, and norm are reallocated per (batch, channel) iteration. On large inputs this can materially erode throughput.

Proposed fix (reuse slice buffer/tensor in-loop)

-        for (int b = 0; b < batchSize; b++)
+        var sliceData = new T[spatialSize];
+        var channelTensor = new Tensor<T>(sliceData, [spatialSize]);
+
+        for (int b = 0; b < batchSize; b++)
         {
             for (int c = 0; c < channels; c++)
             {
                 T mean = NumOps.FromDouble(c < means.Length ? means[c] : 0.5);
                 T std = NumOps.FromDouble(c < stds.Length ? stds[c] : 0.5);

                 int offset = (b * channels + c) * spatialSize;

-                // Extract channel slice as a flat tensor
-                var sliceData = new T[spatialSize];
                 srcSpan.Slice(offset, spatialSize).CopyTo(sliceData);
-                var channelTensor = new Tensor<T>(sliceData, [spatialSize]);

                 // Engine-accelerated: (channel - mean) / std
                 var centered = Engine.TensorSubtractScalar(channelTensor, mean);
                 var norm = Engine.TensorDivideScalar(centered, std);

                 // Copy result back
                 norm.Data.Span.CopyTo(dstSpan.Slice(offset, spatialSize));
             }
         }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Document/PixelToSequence/MATCHA.cs` around lines 626 - 636, The hot loop
allocates sliceData, channelTensor, centered, and norm on every (batch, channel)
iteration; to fix, move a reusable buffer and tensor allocations outside the
loops (e.g., rent a T[] from ArrayPool<T> as sliceBuffer and create a single
Tensor<T> channelTensorReused with shape [spatialSize]), reuse them each
iteration, and prefer in-place or output-writing Engine APIs (or add/consume
overloads that take an output Tensor) instead of creating new centered/norm
tensors; after processing, copy the result from the reused tensor into
dstSpan.Slice(offset, spatialSize) and return the rented buffer when done.

src/NeuralNetworks/Layers/GRULayer.cs (1)

1422-1438: ⚠️ Potential issue | 🔴 Critical

Blocking: invalidate engine/GPU parameter caches after in-place SetParameters.

This now guards the input length, but it still leaves the layer in an inconsistent state after the bulk copy. SetParameters(...) mutates the live CPU tensors without invalidating persistent tensor state or the stacked GPU weights, so subsequent GPU execution can keep using stale parameters.

Suggested fix

 public override void SetParameters(Vector<T> parameters)
 {
     if (parameters.Length != ParameterCount)
         throw new ArgumentException(
             $"Expected {ParameterCount} parameters but got {parameters.Length}.",
             nameof(parameters));

     int idx = 0;
     parameters.Slice(idx, _Wz.Length).AsSpan().CopyTo(_Wz.Data.Span); idx += _Wz.Length;
     parameters.Slice(idx, _Wr.Length).AsSpan().CopyTo(_Wr.Data.Span); idx += _Wr.Length;
     parameters.Slice(idx, _Wh.Length).AsSpan().CopyTo(_Wh.Data.Span); idx += _Wh.Length;
     parameters.Slice(idx, _Uz.Length).AsSpan().CopyTo(_Uz.Data.Span); idx += _Uz.Length;
     parameters.Slice(idx, _Ur.Length).AsSpan().CopyTo(_Ur.Data.Span); idx += _Ur.Length;
     parameters.Slice(idx, _Uh.Length).AsSpan().CopyTo(_Uh.Data.Span); idx += _Uh.Length;
     parameters.Slice(idx, _bz.Length).AsSpan().CopyTo(_bz.Data.Span); idx += _bz.Length;
     parameters.Slice(idx, _br.Length).AsSpan().CopyTo(_br.Data.Span); idx += _br.Length;
     parameters.Slice(idx, _bh.Length).AsSpan().CopyTo(_bh.Data.Span);
+
+    Engine.InvalidatePersistentTensor(_Wz);
+    Engine.InvalidatePersistentTensor(_Wr);
+    Engine.InvalidatePersistentTensor(_Wh);
+    Engine.InvalidatePersistentTensor(_Uz);
+    Engine.InvalidatePersistentTensor(_Ur);
+    Engine.InvalidatePersistentTensor(_Uh);
+    Engine.InvalidatePersistentTensor(_bz);
+    Engine.InvalidatePersistentTensor(_br);
+    Engine.InvalidatePersistentTensor(_bh);
+
+    _gpuWz?.Dispose(); _gpuWz = null;
+    _gpuWr?.Dispose(); _gpuWr = null;
+    _gpuWh?.Dispose(); _gpuWh = null;
+    _gpuUz?.Dispose(); _gpuUz = null;
+    _gpuUr?.Dispose(); _gpuUr = null;
+    _gpuUh?.Dispose(); _gpuUh = null;
+    _gpuBz?.Dispose(); _gpuBz = null;
+    _gpuBr?.Dispose(); _gpuBr = null;
+    _gpuBh?.Dispose(); _gpuBh = null;
+
+    InvalidateGpuStackedWeights();
 }

As per coding guidelines: “Every PR must contain production-ready code.”

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/GRULayer.cs` around lines 1422 - 1438,
SetParameters currently mutates the CPU backing buffers
(_Wz,_Wr,_Wh,_Uz,_Ur,_Uh,_bz,_br,_bh) in-place but does not invalidate any
persistent tensor state or GPU/device caches, leaving GPU execution with stale
weights; after completing the parameter copies in SetParameters, call the
appropriate cache-invalidation API for each affected tensor or for the layer
(e.g., invoke the tensor-level method such as .Invalidate() or .MarkDirty() on
_Wz.Data, _Wr.Data, _Wh.Data, _Uz.Data, _Ur.Data, _Uh.Data, _bz.Data, _br.Data,
_bh.Data, or call the engine-level helper like
Engine.InvalidateParameterCache(this) if available) so that stacked GPU weights
and persistent tensors are refreshed before the next execution.

src/NeuralNetworks/HopeNetwork.cs (1)

163-168: ⚠️ Potential issue | 🟠 Major

Predict() may still be stateful even after these training-mode guards.

These changes stop _associativeMemory and _metaState writes, but Forward() still runs _contextFlow.PropagateContext(...) and every RecurrentLayer<T>.Forward(...). If either of those advances internal state in eval mode, repeated evaluations with identical parameters will still diverge and break the optimizer’s parameter-only cache assumption. Please either snapshot/restore that state in Predict() or make the eval path explicitly read-only.

#!/bin/bash
set -euo pipefail

mapfile -t files < <(fd -i '^(ContextFlow|IContextFlow|RecurrentLayer)\.cs$' src)

printf 'Files to inspect:\n'
printf ' - %s\n' "${files[@]}"

printf '\n=== ContextFlow<T>.PropagateContext(...) ===\n'
rg -n -C4 '\bPropagateContext\s*\(' "${files[@]}"

printf '\n=== RecurrentLayer<T>.Forward(...) ===\n'
rg -n -C6 '\bForward\s*\(' $(fd -i '^RecurrentLayer\.cs$' src)

printf '\n=== Training-mode guards / likely state fields ===\n'
rg -n -C3 '\b(IsTrainingMode|SetTrainingMode|ResetState|_state|_hidden|_context)\b' "${files[@]}"

Expected result: if PropagateContext(...) or RecurrentLayer<T>.Forward(...) assign to hidden/context state without an eval-mode guard, this concern is confirmed.

Also applies to: 199-204, 402-424

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/HopeNetwork.cs` around lines 163 - 168, Predict() is still
potentially stateful because Forward() calls _contextFlow.PropagateContext(...)
and RecurrentLayer<T>.Forward(...) which may mutate internal hidden/context
state even in eval mode; update Predict() to either snapshot and restore any
mutable runtime state (_metaState, _associativeMemory, any _hidden/_context
fields) around the evaluation calls or change the eval path to use read-only
variants of ContextFlow.PropagateContext and RecurrentLayer<T>.Forward (or add
explicit IsTrainingMode guards inside those methods) so no internal state is
modified during prediction; locate usages in Predict(), Forward(),
_contextFlow.PropagateContext, and RecurrentLayer<T>.Forward to implement the
snapshot/restore or make-read-only change.

src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs (1)

567-588: ⚠️ Potential issue | 🔴 Critical

Blocking: noisy gating is numerically unstable and now differs between CPU and GPU training.

Lines 573-575 implement softplus as log(1 + exp(x)), which will overflow for large positive logits and can corrupt routing with Inf/NaN. This behavior was also only added to Forward, so ForwardGpu trains with different routing semantics than CPU. Please move this into one shared helper, use a stable softplus, and call it from both execution paths.

Suggested direction

- // Softplus(x) = log(1 + exp(x)) — noise scale
- var expLogits = Engine.TensorExp(routingLogits);
- var onePlusExp = Engine.TensorAddScalar(expLogits, NumOps.One);
- var noiseScale = Engine.TensorLog(onePlusExp);
+ var noiseScale = ComputeStableNoiseScale(routingLogits);

- if (IsTrainingMode)
- {
-     ...
-     routingLogits = Engine.TensorAdd(routingLogits, Engine.TensorMultiply(noise, noiseScale));
- }
+ if (IsTrainingMode)
+ {
+     routingLogits = ApplyNoisyGating(routingLogits);
+ }

Implement ApplyNoisyGating(...) once and reuse it from both Forward and ForwardGpu.

As per coding guidelines, "All methods have complete, production-ready implementations" and "Simplified implementations: Code that takes shortcuts like hardcoded values instead of proper logic" must be flagged as blocking.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs` around lines 567 - 588,
The current per-branch softplus implementation (in the Forward block using
Engine.TensorExp/Log) overflows for large positive routingLogits and is only
present in Forward (not ForwardGpu), causing CPU/GPU divergence; extract the
noisy-gating logic into a single helper ApplyNoisyGating(...) that both Forward
and ForwardGpu call, and inside it compute softplus in a numerically stable way
(softplus(x) = max(x,0) + log(1 + exp(-|x|))) using available Engine tensor ops
rather than plain log(1+exp(x)), generate the standard-normal noise (Box‑Muller)
there and apply routingLogits = routingLogits + noise * noiseScale so both
execution paths share identical, stable behavior.

src/Diffusion/Conditioning/T5TextConditioner.cs (1)

128-170: ⚠️ Potential issue | 🟠 Major

Make the lazy tensor caches thread-safe.

EnsureWeightTensorsBuilt() and RMSNormEngine() still use unsynchronized null-check initialization. Concurrent EncodeText() calls can race while populating _attnWeightTensors, _ffnWeightTensors, _rmsOnesTensor, and _rmsEpsTensor, which risks duplicate work and publishing partially-built state.

Also applies to: 376-392

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Diffusion/Conditioning/T5TextConditioner.cs` around lines 128 - 170,
EnsureWeightTensorsBuilt and RMSNormEngine perform unsynchronized lazy
initialization causing races; make their caches thread-safe by introducing a
private sync object (e.g., _weightsInitLock) and wrap the null-check +
initialization in a lock and re-check pattern (double-check locking) so that
only one thread builds _attnWeightTensors and _ffnWeightTensors; do the same for
_rmsOnesTensor and _rmsEpsTensor inside RMSNormEngine (use the same or a
separate lock), re-checking the fields inside the lock before populating to
avoid publishing partially-built arrays or duplicate work.

src/Clustering/Density/HDBSCAN.cs (3)

445-453: ⚠️ Potential issue | 🔴 Critical

Birth the cluster when two small components first cross minClusterSize.

This branch still just unions the components. If size1 + size2 >= minClusterSize, that merge is the cluster's birth event; skipping it drops the birth lambda and can leave the condensed tree empty on the terminal merge, returning all noise. Allocate the cluster id here and materialize that birth into the condensed-tree state.

Suggested fix

             if (!big1 && !big2)
             {
-                // Both too small — just merge, no condensed tree entry
                 Union(parent, size, root1, root2);
                 int newRoot = Find(parent, root1);
-                // Propagate cluster label if either had one
                 if (isCluster1) clusterLabel[newRoot] = clusterLabel[root1];
                 else if (isCluster2) clusterLabel[newRoot] = clusterLabel[root2];
+                else if (size[newRoot] >= minClusterSize)
+                {
+                    int newClusterId = nextCluster++;
+                    clusterLabel[newRoot] = newClusterId;
+                    // Also record this birth event so extraction can see the new cluster.
+                }
             }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Clustering/Density/HDBSCAN.cs` around lines 445 - 453, The current branch
in HDBSCAN.cs that handles merging two small components (using Union(parent,
size, root1, root2) and computing newRoot via Find) must also detect if
size[root1] + size[root2] >= minClusterSize and, if so, allocate a new cluster
id and record the birth event into the condensed-tree state instead of silently
merging; modify the block handling (!big1 && !big2) to compute the combined
size, create/assign a cluster id, set clusterLabel[newRoot] appropriately, and
call the existing condensed-tree/materialization routines used elsewhere for
birth events so the birth lambda and cluster entry are produced for that merge.

---

176-181: ⚠️ Potential issue | 🔴 Critical

finalParent still cannot recover embedded core points.

Find(finalParent, i) gives a union-find component root, not a condensed-tree cluster id. That root is typically < n, so both fallback walkers short-circuit before they can climb parentLookup, and the recovered-point path still leaves those points at -1. Return a UF-root→cluster-id mapping from BuildCondensedTree and seed the fallback from that value instead of the raw UF root.

Suggested fix

-        int[] finalParent;
-        (_condensedTree, finalParent) = BuildCondensedTree(mst, n, _options.MinClusterSize);
+        int[] finalParent;
+        int[] finalClusterLabel;
+        (_condensedTree, finalParent, finalClusterLabel) = BuildCondensedTree(mst, n, _options.MinClusterSize);

-        var clusterLabels = ExtractClusters(_condensedTree, n, _options.ClusterSelection, _options.AllowSingleCluster, finalParent);
+        var clusterLabels = ExtractClusters(
+            _condensedTree,
+            n,
+            _options.ClusterSelection,
+            _options.AllowSingleCluster,
+            finalParent,
+            finalClusterLabel);

And in the fallback path:

-                    int root = Find(ufParent, i);
-                    int current = root;
+                    int root = Find(ufParent, i);
+                    int current = finalClusterLabel[root];

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Clustering/Density/HDBSCAN.cs` around lines 176 - 181, BuildCondensedTree
currently returns finalParent (union-find parent array) which Find(finalParent,
i) yields a UF root (< n) and causes the fallback walkers in ExtractClusters to
short-circuit before climbing parentLookup, leaving embedded core points as -1;
modify BuildCondensedTree to also compute and return a UF-root→cluster-id
mapping (e.g., rootToClusterId) alongside _condensedTree and finalParent, then
in ExtractClusters use rootToClusterId[Find(finalParent, i)] as the seeded
cluster id for both fallback walkers (instead of the raw UF root) so the walkers
can correctly climb parentLookup and recover embedded core points.

---

637-675: ⚠️ Potential issue | 🟠 Major

Do not drive the EOM sweep from OrderByDescending(clusterId).

Cluster ids are allocation order, not a post-order traversal of the condensed tree. This can evaluate a parent before one of its descendants has finished propagating stability, so the keep-parent-vs-keep-children decision becomes order-dependent. Build a real descendant-first traversal from children and run the EOM sweep in that order.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Clustering/Density/HDBSCAN.cs` around lines 637 - 675, The EOM sweep
currently iterates clusterList = clusterNodes.OrderByDescending(c => c) which
uses allocation ids rather than a descendant-first (post-order) traversal; this
can evaluate a parent before its descendants finish propagating stability.
Replace the OrderByDescending-based ordering with a true descendant-first
traversal built from the children dictionary (e.g., perform a post-order DFS
starting from root/top-level nodes that collects cluster ids after visiting
children) and then iterate that post-order list in the existing sweep code
(keeping references to stability, isCluster, children, clusterNodes and calling
MarkDescendantsNotCluster as before) so parents are processed only after all
descendants have propagated their stability.

src/NeuralNetworks/Layers/EmbeddingLayer.cs (1)

368-371: ⚠️ Potential issue | 🟠 Major

Redundant embedding allocation in initialization path.

Line 370 recreates _embeddingTensor even though it is already allocated in the constructor (Line 299). This adds an avoidable full-tensor allocation and conflicts with the “in-place” intent in the comment.

💡 Suggested fix

- _embeddingTensor = new Tensor<T>([vocabularySize, embeddingDimension]);
+ _embeddingTensor = Tensor<T>.CreateRandom(vocabularySize, embeddingDimension);
  InitializeParameters();

- _embeddingTensor = Tensor<T>.CreateRandom(vocabSize, embeddingDim);
  var halfTensor = TensorAllocator.Rent<T>([vocabSize, embeddingDim]);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/EmbeddingLayer.cs` around lines 368 - 371, The
initialization mistakenly re-allocates _embeddingTensor with
Tensor<T>.CreateRandom even though _embeddingTensor was already allocated in the
constructor; remove that redundant allocation and perform the randomization
in-place into the existing _embeddingTensor (or into the rented buffer and then
copy/assign into _embeddingTensor as intended) using the TensorAllocator buffer
and in-place ops so you avoid creating a second full-size tensor—update the
block that currently calls Tensor<T>.CreateRandom and instead call the in-place
random/fill routine on _embeddingTensor (or fill the rented tensor and move it
into _embeddingTensor) and ensure the rented halfTensor is used for the
shift/scale steps and returned via TensorAllocator.Return.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 35dcc5d8-b910-4ec4-8cba-e9eda7731090

📥 Commits

Reviewing files that changed from the base of the PR and between 2c92465 and cc4b7f2.

📒 Files selected for processing (48)

• coverlet.runsettings
• src/ActivationFunctions/BinarySpikingActivation.cs
• src/ActivationFunctions/LiSHTActivation.cs
• src/ActivationFunctions/RReLUActivation.cs
• src/ActivationFunctions/ScaledTanhActivation.cs
• src/ActivationFunctions/SphericalSoftmaxActivation.cs
• src/Audio/Enhancement/CMGAN.cs
• src/Audio/Enhancement/SpikingFullSubNet.cs
• src/Augmentation/Image/ImageTensor.cs
• src/Autodiff/Testing/TensorOperationsVerification.cs
• src/Classification/Boosting/NGBoostClassifier.cs
• src/Clustering/Density/HDBSCAN.cs
• src/Data/Audio/AudioFileDataset.cs
• src/Data/Formats/ParquetDataLoader.cs
• src/Data/Geometry/KittiDataLoader.cs
• src/Data/Geometry/NuScenesDataLoader.cs
• src/Data/Geometry/PointCloudDatasetLoaderBase.cs
• src/Data/Geometry/SemanticKittiDataLoader.cs
• src/Data/Geometry/WaymoDataLoader.cs
• src/Data/Graph/ProteinDataLoader.cs
• src/Data/Graph/TemporalGraphDataLoader.cs
• src/Data/Graph/Wikidata5mDataLoader.cs
• src/Data/Loaders/InMemoryDataLoader.cs
• src/Data/Text/StreamingTextDataset.cs
• src/Data/Text/TokenizedTextDataset.cs
• src/Data/Video/VideoFrameDataset.cs
• src/Data/Vision/ImageClassificationDataset.cs
• src/Data/Vision/ImageFolderDataset.cs
• src/Diffusion/Conditioning/T5TextConditioner.cs
• src/Diffusion/Control/ControlNetModel.cs
• src/Document/PixelToSequence/MATCHA.cs
• src/Finance/Data/FinancialDataLoader.cs
• src/FitnessCalculators/CompressionAwareFitnessCalculator.cs
• src/GaussianProcesses/DeepGaussianProcess.cs
• src/Helpers/FeatureSelectorHelper.cs
• src/Interfaces/IFitnessCalculator.cs
• src/Models/Options/NBEATSModelOptions.cs
• src/NeuralNetworks/HopeNetwork.cs
• src/NeuralNetworks/Layers/CapsuleLayer.cs
• src/NeuralNetworks/Layers/EmbeddingLayer.cs
• src/NeuralNetworks/Layers/GRULayer.cs
• src/NeuralNetworks/Layers/GraphSAGELayer.cs
• src/NeuralNetworks/Layers/GraphTransformerLayer.cs
• src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs
• src/NeuralNetworks/Layers/MaskingLayer.cs
• src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs
• src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs
• src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs

[coderabbitai]

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/AiModelBuilder.cs`:
- Around line 2502-2505: The preprocessing pipeline was fitted earlier on XTrain
but you later call _preprocessingPipeline.Transform(clusterX), causing
normalization to use training-split statistics; instead refit (or FitTransform)
the pipeline on the actual clustering dataset before transforming so the
full-data branch uses clusterX statistics — i.e., when preparing fullX for
clustering, call the pipeline's Fit (or FitTransform) on clusterX (using the
same _preprocessingPipeline instance or a new one) and then set fullX to the
transformed result (and leave fullY = clusterY unchanged).
- Around line 2482-2500: The clustering path currently calls
_dataPreparationPipeline.FitResample / FitResampleTensor and thus uses clusterY
(labels), which leaks labels into clustering; change this so clustering only
sees X-only transforms: before calling FitResample/FitResampleTensor, inspect
_dataPreparationPipeline and either (a) build a filtered pipeline that excludes
label-aware steps (e.g., resamplers/augmenters) and only run Fit/Transform on X,
or (b) if no safe filtered pipeline API exists, skip
FitResample/FitResampleTensor entirely and only apply non-label transforms (call
their Transform on preparedX) so you never pass clusterY into
_dataPreparationPipeline; update the code around clusterX/clusterY,
preparedX/preparedY, and the _dataPreparationPipeline usage to use the filtered
pipeline or X-only transforms for clustering.

In `@src/Helpers/LayerHelper.cs`:
- Around line 31254-31264: The visual encoder is being appended to the same
sequential stream as the audio encoder (you yield a second
DenseLayer<T>(inputSize, embeddingDimension, ...) after audio outputs
embeddingDimension), causing shape/flow errors; refactor so audio and visual are
built as separate pipelines instead of one IEnumerable<ILayer<T>> chain—either
return a structure with two IEnumerables or construct a branching/composite
layer (e.g. a Fork/Parallel/Composite layer) that contains two independent
sequences (audio: DenseLayer<T>(inputSize,...) -> N x MultiHeadAttentionLayer<T>
-> DenseLayer<T>(embeddingDimension,...); visual: DenseLayer<T>(inputSize,...)
-> N x MultiHeadAttentionLayer<T> -> DenseLayer<T>(embeddingDimension,...)) and
ensure the composed layer accepts separate inputs so DenseLayer<T> sizes chain
correctly.
- Around line 21489-21494: Update the XML doc for the public factory method
CreateDefaultVoiceCraftLayers to include full contract tags: add a <summary>
that describes the method, add <param> entries for the architecture parameter
(NeuralNetworkArchitecture<T>) and all other parameters (hiddenDim, numLayers,
numHeads, codebookSize, dropoutRate), add a <returns> describing the
IEnumerable<ILayer<T>> result, and add a <remarks> block that contains a
<para><b>For Beginners:</b> explanation of what the generated layers represent
and when to use this factory; ensure the architecture parameter is documented to
explain expected topology/constraints. Ensure the XML structure matches the
other LayerHelper CreateDefault{ModelName}Layers patterns in the file.
- Around line 32704-32730: The parameter numResidualLayers is unused and the
code always yields exactly two GRULayer<T> instances; replace the two hardcoded
yields with a loop that yields numResidualLayers GRULayer<T> instances so
callers can control the recurrent stack: iterate i from 0 to numResidualLayers-1
and yield return new GRULayer<T>(hiddenDim, hiddenDim, false,
(IActivationFunction<T>?)null) for each iteration (keeping the default
numResidualLayers = 2 intact).
- Around line 4621-4627: The model head currently returns raw logits via
DenseLayer<T> with IdentityActivation<T>, but the loss implementations
CrossEntropyLoss<T> and BinaryCrossEntropyLoss<T> expect probabilities; either
restore a probability-producing head (replace IdentityActivation<T> with
SoftmaxActivation<T> for multi-class and SigmoidActivation<T> for binary in the
DenseLayer<T> returned by the DenseNet head) or else update the loss
implementations to accept logits; make the quick fix by changing the activation
passed to the DenseLayer<T> in the DenseNet head (or alternately update
CrossEntropyLoss<T>/BinaryCrossEntropyLoss<T> to apply LogSoftmax/Sigmoid
internally) so model outputs and losses agree.

In `@src/NeuralNetworks/Layers/RBFLayer.cs`:
- Around line 61-62: The TrainableParameter registration for fields annotated
with [TrainableParameter] (_centers and _widths) must not be left pointing at
stale Tensor<T> instances when SetParameters() or UpdateParameters() rebind
those fields; locate where SetParameters()/UpdateParameters() currently replace
_centers/_widths and instead update the existing Tensor<T> objects in-place
(e.g., copy/assign the new data into the existing tensor buffer via the tensor
API such as CopyFrom/SetData/AssignFrom or similar) so the registered identity
remains stable; if the tensor API doesn't allow in-place update, then after
replacing the field call the parameter registry refresh routine (the code path
that registers TrainableParameter entries) to re-register the new tensor
instance immediately.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 7b680873-b425-40c0-854a-9afd397fd1f6

📥 Commits

Reviewing files that changed from the base of the PR and between cc4b7f2 and 2b3f6cc.

📒 Files selected for processing (6)

• .github/scripts/report-slow-tests.ps1
• Directory.Packages.props
• src/AiModelBuilder.cs
• src/Audio/Generation/VoiceCraft.cs
• src/Helpers/LayerHelper.cs
• src/NeuralNetworks/Layers/RBFLayer.cs

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of files (300). Try reducing the number of changed files and requesting a review from Copilot again.

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (4)

src/NeuralNetworks/Layers/RBFLayer.cs (1)

329-340: ⚠️ Potential issue | 🔴 Critical

BLOCKING: UpdateParameters still replaces tensor references, breaking registered parameter identity.

While SetParameters (lines 415-417) was correctly fixed to use in-place span copies, UpdateParameters still reassigns _centers and _widths to new Tensor<T> objects returned by Engine.TensorSubtract. After this method runs, the trainable parameter registry will hold stale references to the old tensors.

This directly contradicts the PR's goal of preserving registered parameter tensor references and will cause the same drift issue flagged in past reviews.

🔧 Proposed fix: Update tensors in-place

 public override void UpdateParameters(T learningRate)
 {
     if (_centersGradient == null || _widthsGradient == null)
         throw new InvalidOperationException("Backward pass must be called before updating parameters.");

-    // Use Engine.TensorSubtract and TensorMultiplyScalar for GPU/CPU acceleration
-    var scaledCentersGradient = Engine.TensorMultiplyScalar(_centersGradient, learningRate);
-    _centers = Engine.TensorSubtract(_centers, scaledCentersGradient);
-
-    var scaledWidthsGradient = Engine.TensorMultiplyScalar(_widthsGradient, learningRate);
-    _widths = Engine.TensorSubtract(_widths, scaledWidthsGradient);
+    // Compute updates
+    var scaledCentersGradient = Engine.TensorMultiplyScalar(_centersGradient, learningRate);
+    var newCenters = Engine.TensorSubtract(_centers, scaledCentersGradient);
+    
+    var scaledWidthsGradient = Engine.TensorMultiplyScalar(_widthsGradient, learningRate);
+    var newWidths = Engine.TensorSubtract(_widths, scaledWidthsGradient);
+    
+    // Write in-place to preserve registered parameter tensor references
+    newCenters.Data.Span.CopyTo(_centers.Data.Span);
+    newWidths.Data.Span.CopyTo(_widths.Data.Span);
 }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/RBFLayer.cs` around lines 329 - 340, Update
UpdateParameters so it does not replace the _centers and _widths references;
currently it reassigns them to the results of Engine.TensorSubtract which breaks
the registered parameter identity. Instead compute the scaled gradients from
_centersGradient and _widthsGradient and apply those adjustments in-place to the
existing tensors ( _centers and _widths ) using the Engine's in-place ops or
element-wise copy/subtract helpers (e.g., TensorMultiplyScalarInto /
TensorSubtractInPlace or equivalent) so the original Tensor<T> objects remain
the same while their values are updated.

src/NeuralNetworks/Layers/DenseLayer.cs (3)

1423-1445: ⚠️ Potential issue | 🟡 Minor

Velocity tensors should also be invalidated in Dispose().

The momentum velocity tensors (_weightsVelocity, _biasesVelocity) are registered with the GPU engine as PersistentTensorRole.OptimizerState (lines 1140, 1146), but they are not invalidated in Dispose(). While they are created with new Tensor<T>() (not rented), they should still be invalidated to release any GPU handles.

🔧 Proposed addition to Dispose()

 protected override void Dispose(bool disposing)
 {
     if (disposing)
     {
         // Release GPU handles for persistent tensors
         Engine.InvalidatePersistentTensor(_weights);
         Engine.InvalidatePersistentTensor(_biases);
+        
+        // Release GPU handles for optimizer state tensors if present
+        if (_weightsVelocity != null)
+            Engine.InvalidatePersistentTensor(_weightsVelocity);
+        if (_biasesVelocity != null)
+            Engine.InvalidatePersistentTensor(_biasesVelocity);
+        _weightsVelocity = null;
+        _biasesVelocity = null;

         // Clear other managed resources (CPU)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/DenseLayer.cs` around lines 1423 - 1445, Dispose
currently invalidates persistent GPU tensors but omits the optimizer momentum
tensors; update the Dispose(bool disposing) implementation to call
Engine.InvalidatePersistentTensor for both _weightsVelocity and _biasesVelocity
when disposing is true so their GPU handles are released, and also set
_weightsVelocity and _biasesVelocity to null alongside the other managed fields
to avoid dangling references.

---

1038-1075: ⚠️ Potential issue | 🔴 Critical

BLOCKING: Memory leak when resizing weights — old tensor not returned to pool.

When actualInputSize differs from the current weight shape, a new tensor is rented (line 1048) and assigned to _weights (line 1072). The previously rented _weights tensor is orphaned without being returned to the pool.

This leak accumulates each time input dimensions change dynamically.

🐛 Proposed fix to return old tensor before replacement

 private void EnsureWeightShapeForInput(int actualInputSize)
 {
     // Weights are [inputSize, outputSize]
     if (_weights.Shape[0] == actualInputSize)
     {
         return;
     }

     int existingInputSize = _weights.Shape[0];
     int outputSize = _weights.Shape[1];
+    var oldWeights = _weights;
     var resizedWeights = TensorAllocator.Rent<T>([actualInputSize, outputSize]);

     int sharedInputSize = Math.Min(existingInputSize, actualInputSize);
     for (int i = 0; i < sharedInputSize; i++)
     {
         for (int o = 0; o < outputSize; o++)
         {
             resizedWeights[i, o] = _weights[i, o];
         }
     }

     if (actualInputSize > sharedInputSize)
     {
         T scale = NumOps.FromDouble(Math.Sqrt(2.0 / (actualInputSize + outputSize)));
         var random = RandomHelper.CreateSecureRandom();
         for (int i = sharedInputSize; i < actualInputSize; i++)
         {
             for (int o = 0; o < outputSize; o++)
             {
                 resizedWeights[i, o] = NumOps.Multiply(scale, NumOps.FromDouble(random.NextDouble() * 2 - 1));
             }
         }
     }

+    // Return old tensor to pool before replacing
+    Engine.InvalidatePersistentTensor(oldWeights);
+    TensorAllocator.Return(oldWeights);
+
     _weights = resizedWeights;
+    RegisterTrainableParameter(_weights, PersistentTensorRole.Weights);
     _weightsGradient = null;
     UpdateInputShape([actualInputSize]);
 }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/DenseLayer.cs` around lines 1038 - 1075, In
EnsureWeightShapeForInput, when resizing weights you rent a new tensor
(resizedWeights) and then overwrite _weights without returning the old tensor to
the allocator; before assigning _weights = resizedWeights you must return the
previous _weights to the pool (e.g., call TensorAllocator.Return/Dispose on the
old tensor if non-null) and also clear/return _weightsGradient if it references
the old buffer; then assign _weights = resizedWeights, set _weightsGradient =
null, and call UpdateInputShape(actualInputSize) as before to avoid leaking
rented tensors.

---

713-724: ⚠️ Potential issue | 🟠 Major

Resource leak: Old weights tensor not returned when replaced via SetWeights().

When SetWeights() replaces _weights with the provided tensor, the previously rented tensor is not returned to the pool. While this method is protected and likely called less frequently than EnsureWeightShapeForInput, it still creates a leak when used (e.g., for transfer learning or loading pre-trained models).

🔧 Proposed fix

 protected override void SetWeights(Tensor<T> weights)
 {
     if (weights == null)
     {
         throw new ArgumentNullException(nameof(weights));
     }

     // Ensure weights are initialized before validation (supports lazy initialization)
     EnsureInitialized();

     // Validate dimensions against current weights: [inputSize, outputSize]
     if (weights.Shape[0] != _weights.Shape[0] || weights.Shape[1] != _weights.Shape[1])
     {
         throw new ArgumentException(
             $"Weight tensor dimensions must be {_weights.Shape[0]}x{_weights.Shape[1]}, but got {weights.Shape[0]}x{weights.Shape[1]}");
     }

+    // Return old rented tensor to pool
+    var oldWeights = _weights;
+    Engine.InvalidatePersistentTensor(oldWeights);
+    TensorAllocator.Return(oldWeights);
+
     // Set the weights directly
     _weights = weights;
+    RegisterTrainableParameter(_weights, PersistentTensorRole.Weights);

     // Update input shape if needed - Shape[0] is inputSize in new convention

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/DenseLayer.cs` around lines 713 - 724, SetWeights
currently overwrites the _weights tensor without returning or disposing the
previous tensor, leaking rented GPU memory; before assigning _weights = weights
in the SetWeights method, check if the existing _weights is non-null and not the
same instance as the incoming weights and return it to the pool or dispose it
(e.g., call Engine.ReturnTensor(_weights) if your engine exposes that, or
_weights.Dispose()), then assign the new tensor and call
Engine.InvalidatePersistentTensor(_weights) as before.

♻️ Duplicate comments (2)

src/NeuralNetworks/Layers/ReservoirLayer.cs (1)

330-335: ⚠️ Potential issue | 🔴 Critical

Rebuilding outputs via raw span copies detaches the returned tensor from autodiff.

Lines 330-335 materialize a fresh Tensor<T> and copy row data into it with Data.Span.CopyTo. That bypasses the engine/tape, so gradients from the returned batch output stop here even though the per-step state updates were engine-backed. Keep the final batch assembly on engine/tape-aware ops as well.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/Layers/ReservoirLayer.cs` around lines 330 - 335, The
current assembly in ReservoirLayer (where outputs is created with new
Tensor<T>([flatBatch, _reservoirSize]) and rows are copied via Data.Span.CopyTo)
breaks autodiff because it materializes a fresh tensor outside the engine/tape;
replace this manual span-copy with an engine/tape-aware batched assemble
operation (e.g. use the framework's Tensor.Stack/Concat/Gather equivalent) to
combine outputRows into a single tensor along the batch axis so gradients flow
through the returned outputs; operate on outputRows (the engine-backed per-step
tensors) and produce outputs via the engine API rather than allocating and
copying into a raw Tensor<T>.

src/ActivationFunctions/RReLUActivation.cs (1)

148-149: ⚠️ Potential issue | 🟠 Major

Don't reset the training slopes to the inference midpoint.

After sampling per-element alpha values for the tensor forward pass, Line 149 overwrites _alpha with the bounds midpoint. That makes Derivative(T) report a slope that was never used in the forward pass, and any fallback tensor derivative path that still delegates to the scalar derivative will return wrong gradients for negative elements.

Suggested fix direction

-        // Store midpoint as representative alpha for derivative
-        _alpha = NumOps.Divide(NumOps.Add(_lowerBound, _upperBound), NumOps.FromDouble(2.0));
+        // Preserve the actual sampled slopes for any non-graph derivative path.
+        _lastAlphaTensor = alphaTensor;

You would then want the tensor-derivative path to read _lastAlphaTensor instead of collapsing back to a scalar midpoint.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ActivationFunctions/RReLUActivation.cs` around lines 148 - 149, The code
in RReLUActivation resets the training slopes by overwriting _alpha with the
midpoint, causing Derivative(T) to report a slope that wasn't used in the
forward tensor pass; remove the assignment that collapses per-element sampled
alphas into the scalar midpoint and ensure any tensor-derivative path reads
per-element slopes from _lastAlphaTensor (and only use a scalar midpoint for
pure inference paths), updating references in methods like Derivative(T) and the
forward/backward code to use _lastAlphaTensor when available instead of _alpha.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/Diffusion/NoisePredictors/FluxDoubleStreamPredictor.cs`:
- Around line 150-156: Replace the per-call List<Vector<T>> allocation with a
pre-sized Vector<T>[]: compute the exact component count (1 for _patchEmbed +
_doubleBlocks.Count + _singleBlocks.Count + 1 for _finalLayer), allocate
Vector<T>[] arr = new Vector<T>[count], fill arr[0] =
_patchEmbed.GetParameters(), then loop to populate double/single blocks and set
the final layer, and call Vector<T>.Concatenate(arr) (do the same replacement
for the similar code path that currently allocates a List and calls ToArray
around Vector<T>.Concatenate in the second occurrence). Ensure you reference and
populate the same methods: _patchEmbed.GetParameters(), each b.GetParameters()
on _doubleBlocks/_singleBlocks, and _finalLayer.GetParameters().
- Around line 179-183: The SetParams method currently slices parameters without
validating total consumption; update SetParams (involving DenseLayer<T> and
SetParameters) to validate that parameters.Length equals offset +
layer.ParameterCount and throw a clear exception (e.g., ArgumentException) if it
does not, so oversized or undersized parameter vectors are rejected rather than
silently ignored.

In `@src/NeuralNetworks/Layers/DenseLayer.cs`:
- Around line 365-368: The Dispose() path for DenseLayer fails to return rented
tensors, causing a pool leak: locate the Dispose() method in DenseLayer (where
Engine.InvalidatePersistentTensor(...) is called) and after invalidating
persistent tensors ensure you call TensorAllocator.Return(_weights) and
TensorAllocator.Return(_biases) (or the appropriate
TensorAllocator.Return<T>(...) overload) and null out those fields; keep the
InvalidatePersistentTensor calls but add the matching TensorAllocator.Return
invocations for _weights and _biases to properly release pooled resources.

In `@src/NeuralNetworks/Layers/ReservoirLayer.cs`:
- Around line 296-298: The code reshapes and transposes _inputWeights inside
ReservoirLayer.Forward on every call (via Engine.Reshape and
Engine.TensorTranspose producing inputWeightsT); instead, add a cached field
(e.g. _inputWeightsTransposed) computed once when _inputWeights is set or when
parameters change and reuse it in Forward (use inputWeightsT =
_inputWeightsTransposed), and invalidate/recompute this cached transposed tensor
only when _inputWeights is replaced or updated so you eliminate per-evaluation
reshape/transpose work.

---

Outside diff comments:
In `@src/NeuralNetworks/Layers/DenseLayer.cs`:
- Around line 1423-1445: Dispose currently invalidates persistent GPU tensors
but omits the optimizer momentum tensors; update the Dispose(bool disposing)
implementation to call Engine.InvalidatePersistentTensor for both
_weightsVelocity and _biasesVelocity when disposing is true so their GPU handles
are released, and also set _weightsVelocity and _biasesVelocity to null
alongside the other managed fields to avoid dangling references.
- Around line 1038-1075: In EnsureWeightShapeForInput, when resizing weights you
rent a new tensor (resizedWeights) and then overwrite _weights without returning
the old tensor to the allocator; before assigning _weights = resizedWeights you
must return the previous _weights to the pool (e.g., call
TensorAllocator.Return/Dispose on the old tensor if non-null) and also
clear/return _weightsGradient if it references the old buffer; then assign
_weights = resizedWeights, set _weightsGradient = null, and call
UpdateInputShape(actualInputSize) as before to avoid leaking rented tensors.
- Around line 713-724: SetWeights currently overwrites the _weights tensor
without returning or disposing the previous tensor, leaking rented GPU memory;
before assigning _weights = weights in the SetWeights method, check if the
existing _weights is non-null and not the same instance as the incoming weights
and return it to the pool or dispose it (e.g., call
Engine.ReturnTensor(_weights) if your engine exposes that, or
_weights.Dispose()), then assign the new tensor and call
Engine.InvalidatePersistentTensor(_weights) as before.

In `@src/NeuralNetworks/Layers/RBFLayer.cs`:
- Around line 329-340: Update UpdateParameters so it does not replace the
_centers and _widths references; currently it reassigns them to the results of
Engine.TensorSubtract which breaks the registered parameter identity. Instead
compute the scaled gradients from _centersGradient and _widthsGradient and apply
those adjustments in-place to the existing tensors ( _centers and _widths )
using the Engine's in-place ops or element-wise copy/subtract helpers (e.g.,
TensorMultiplyScalarInto / TensorSubtractInPlace or equivalent) so the original
Tensor<T> objects remain the same while their values are updated.

---

Duplicate comments:
In `@src/ActivationFunctions/RReLUActivation.cs`:
- Around line 148-149: The code in RReLUActivation resets the training slopes by
overwriting _alpha with the midpoint, causing Derivative(T) to report a slope
that wasn't used in the forward tensor pass; remove the assignment that
collapses per-element sampled alphas into the scalar midpoint and ensure any
tensor-derivative path reads per-element slopes from _lastAlphaTensor (and only
use a scalar midpoint for pure inference paths), updating references in methods
like Derivative(T) and the forward/backward code to use _lastAlphaTensor when
available instead of _alpha.

In `@src/NeuralNetworks/Layers/ReservoirLayer.cs`:
- Around line 330-335: The current assembly in ReservoirLayer (where outputs is
created with new Tensor<T>([flatBatch, _reservoirSize]) and rows are copied via
Data.Span.CopyTo) breaks autodiff because it materializes a fresh tensor outside
the engine/tape; replace this manual span-copy with an engine/tape-aware batched
assemble operation (e.g. use the framework's Tensor.Stack/Concat/Gather
equivalent) to combine outputRows into a single tensor along the batch axis so
gradients flow through the returned outputs; operate on outputRows (the
engine-backed per-step tensors) and produce outputs via the engine API rather
than allocating and copying into a raw Tensor<T>.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: bbdb727c-6763-4d35-a0c8-84f940f5e98a

📥 Commits

Reviewing files that changed from the base of the PR and between 2b3f6cc and 4b56523.

📒 Files selected for processing (5)

• src/ActivationFunctions/RReLUActivation.cs
• src/Diffusion/NoisePredictors/FluxDoubleStreamPredictor.cs
• src/NeuralNetworks/Layers/DenseLayer.cs
• src/NeuralNetworks/Layers/RBFLayer.cs
• src/NeuralNetworks/Layers/ReservoirLayer.cs

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of files (300). Try reducing the number of changed files and requesting a review from Copilot again.

[coderabbitai]

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/AiModelBuilder.cs`:
- Line 2472: The current check sets useFullData from the possibly-wrapped
variable model, but model may be a DDPModel/FSDPModel wrapper so the
Clustering.Base.ClusteringBase<T> test returns false; instead unwrap the
underlying model (follow wrapper types DDPModel and FSDPModel to their
inner/base model) before evaluating the predicate and set useFullData =
innerModel is Clustering.Base.ClusteringBase<T>; implement a small unwrap loop
or helper that peels DDPModel/FSDPModel wrappers to retrieve the real model and
use that for the type check (referencing useFullData, model, DDPModel,
FSDPModel, and Clustering.Base.ClusteringBase<T>).
- Around line 2478-2506: The run metadata currently logs
training/validation/test counts from the original splits even when clustering
uses the full dataset; update the logging to use the effective datasets chosen
for the "direct" path (directX/directY) so counts reflect what was actually
used. Specifically, after the branch that sets directX and directY (and/or
before logging experiment metadata), compute effectiveTrainingSamples from
directX via ConversionsHelper.GetSampleCount<T, TInput>(directX) and set
effectiveValidationSamples and effectiveTestSamples to 0 when useFullData is
true (otherwise derive them from XVal/XTest), then use these effective counts in
the experiment/training-monitor logging instead of the earlier split-based
counts; adjust code paths around useFullData, _preprocessingPipeline,
XTrain/yTrain, XVal/XTest to ensure the values used for logging match the
dataset passed to clustering/training.

In `@src/LossFunctions/CrossEntropyWithLogitsLoss.cs`:
- Around line 128-131: The tape-loss path in ComputeTapeLoss() currently
averages over all axes (using allAxes with Engine.ReduceMean on product), which
divides by the number of classes and mismatches CalculateLoss(); change it to
sum across the logits/class axis first (reduce/sum over the class/logits
dimension of the product tensor) and then compute the mean only across the
batch/sample axes (exclude the class axis when calling ReduceMean), so the
result equals -Σ target_i * log_softmax_i per sample averaged over samples;
locate ComputeTapeLoss(), the product variable, and the Engine.ReduceMean(...)
call to apply this change.

In `@src/NeuralNetworks/DenseNetNetwork.cs`:
- Around line 146-149: The current switch in DenseNetNetwork.cs incorrectly maps
NeuralNetworkTaskType.BinaryClassification and MultiLabelClassification to
CrossEntropyWithLogitsLoss<T>, which is for mutually-exclusive softmax targets;
change the mapping so BinaryClassification and MultiLabelClassification do not
use CrossEntropyWithLogitsLoss<T> — instead call
NeuralNetworkHelper<T>.GetDefaultLossFunction(taskType) (or the appropriate
binary/multi-label loss provided by NeuralNetworkHelper) for those task types so
binary (including outputSize==1) and independent-label semantics remain correct;
keep CrossEntropyWithLogitsLoss<T> only for true mutually-exclusive multi-class
targets.

In `@src/NeuralNetworks/Layers/ParallelStreamsLayer.cs`:
- Around line 239-260: SetParameters currently slices the incoming Vector<T>
without verifying its length against the layer's ParameterCount; add a guard at
the start of the SetParameters method to validate that parameters.Length (or
Count) equals this.ParameterCount and throw an ArgumentException (or
ArgumentOutOfRangeException) with a clear message if not; then proceed to
iterate over _streamA and _streamB and call
layer.SetParameters(parameters.GetSubVector(offset, count)) as before—this
enforces the XML contract and prevents silently ignored/truncated values.
- Around line 113-130: The ParallelStreamsLayer constructor currently
dereferences streamALayers/streamBLayers and accepts invalid inputSize; add
fail-fast validation at the start of the ParallelStreamsLayer constructor:
validate inputSize > 0 and even (throw
ArgumentOutOfRangeException/ArgumentException with parameter name "inputSize"),
validate streamALayers and streamBLayers are not null (throw
ArgumentNullException for each), and ensure streamAOutputSize/streamBOutputSize
are > 0 (throw ArgumentOutOfRangeException with their parameter names) before
computing _splitSize or creating _streamA/_streamB; keep using _splitSize,
_streamA, _streamB and RegisterSubLayer as-is after these checks.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 599fa050-9acc-44bd-93ad-6749b34b93bf

📥 Commits

Reviewing files that changed from the base of the PR and between 4b56523 and 57791f2.

📒 Files selected for processing (8)

• src/AiModelBuilder.cs
• src/Enums/LossTestInputFormat.cs
• src/FitnessCalculators/CompressionAwareFitnessCalculator.cs
• src/Helpers/LayerHelper.cs
• src/LossFunctions/CrossEntropyWithLogitsLoss.cs
• src/NeuralNetworks/DenseNetNetwork.cs
• src/NeuralNetworks/Layers/CapsuleLayer.cs
• src/NeuralNetworks/Layers/ParallelStreamsLayer.cs

[Copilot]

Pull request overview

Copilot reviewed 67 out of 67 changed files in this pull request and generated 23 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.