fix(densenet): stop holding a forward's activations after it returns

src/NeuralNetworks/Graph/LayerForwardObserver.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using System.Threading;
@@ -216,6 +216,9 @@ public bool HasFanOut
         }
     }
 
+    /// <summary>Drops the recorded calls and the tensors they reference, once they have been read.</summary>
+    public void ReleaseRecording() => _calls.Clear();
+
     /// <inheritdoc />
     public void Dispose()
     {

src/NeuralNetworks/Layers/BatchNormalizationLayer.cs

@@ -869,10 +869,10 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
         // that's never reached when training goes through the tape (tape
         // holds its own intermediate refs already). Skip the assignment
         // when a tape is active so this field doesn't double-root the
-        // input activation; null it out so the previous step's tensor is
-        // eligible for collection.
+        // input activation, and in inference (ShouldCacheForBackward is false
+        // in eval mode), where it held every BN input of the last forward.
         bool tapeActive = AiDotNet.Autodiff.GradientRecording.IsRecording<T>();
-        _lastInput = tapeActive ? null : input;
+        _lastInput = tapeActive || !ShouldCacheForBackward ? null : input;
 
         // Batch statistics are taken over every value a feature has in this batch: the batch
         // axis AND, for convolutional (channels-first) inputs, every spatial position. That is

src/NeuralNetworks/Layers/DenseBlock.cs

@@ -123,7 +123,6 @@ public partial class DenseBlock<T> : LayerBase<T>, ILayerSerializationExtras<T>,
     // Non-readonly: lazy ctor leaves _inputChannels = -1 until
     // OnFirstForward resolves it from the runtime input tensor.
     private int _inputChannels;
-    private List<Tensor<T>>? _layerOutputs;
 
     public override bool SupportsTraining => true;
 
@@ -261,16 +260,13 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
         // before the inner-layer Forwards consume their weights.
         if (!IsShapeResolved) OnFirstForward(input);
 
-        _layerOutputs = new List<Tensor<T>>(_numLayers + 1) { input };
-
         // Current feature maps (accumulated)
         var currentFeatures = input;
 
         foreach (var layer in _layers)
         {
             // Each layer takes ALL previous features as input
             var layerOutput = layer.Forward(currentFeatures);
-            _layerOutputs.Add(layerOutput);
 
             // Concatenate new features with existing features along channel dimension
             currentFeatures = ConcatenateChannels(currentFeatures, layerOutput);
@@ -322,7 +318,6 @@ internal override Dictionary<string, string> GetMetadata()
     /// </summary>
     public override void ResetState()
     {
-        _layerOutputs = null;
         foreach (var layer in _layers)
         {
             layer.ResetState();

src/NeuralNetworks/Layers/DenseBlockLayer.cs

@@ -74,13 +74,6 @@ public partial class DenseBlockLayer<T> : LayerBase<T>, ILayerSerializationExtra
 
     [Scratch]
     private Tensor<T>? _lastInput;
-    private Tensor<T>? _bn1Out;
-    private Tensor<T>? _relu1Out;
-    private Tensor<T>? _conv1Out;
-    [AiDotNet.Attributes.Scratch]
-    private Tensor<T>? _bn2Out;
-    [AiDotNet.Attributes.Scratch]
-    private Tensor<T>? _relu2Out;
 
     public override bool SupportsTraining => true;
 
@@ -203,17 +196,19 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
         var x = input.Shape.Length == 3 ? Engine.Reshape(input, [1, input.Shape[0], input.Shape[1], input.Shape[2]]) : input;
 
         // BN-ReLU-Conv1x1 (DenseNet paper: pre-activation bottleneck)
-        _bn1Out = _bn1.Forward(x);
+        // The intermediates are locals: the tape keeps what backward needs, and a field would hold
+        // every dense layer's activations after the forward, in inference too (DenseNet-121 b16).
+        var bn1Out = _bn1.Forward(x);
         // Scalar activation objects operate eagerly and therefore detach the value from the
         // active GradientTape. Keep the DenseNet pre-activation on the engine graph so gradients
         // reach BN1 and every layer before this block.
-        _relu1Out = Engine.ReLU(_bn1Out);
-        _conv1Out = _conv1x1.Forward(_relu1Out);
+        var relu1Out = Engine.ReLU(bn1Out);
+        var conv1Out = _conv1x1.Forward(relu1Out);
 
         // BN-ReLU-Conv3x3
-        _bn2Out = _bn2.Forward(_conv1Out);
-        _relu2Out = Engine.ReLU(_bn2Out);
-        var output = _conv3x3.Forward(_relu2Out);
+        var bn2Out = _bn2.Forward(conv1Out);
+        var relu2Out = Engine.ReLU(bn2Out);
+        var output = _conv3x3.Forward(relu2Out);
 
         // Remove batch dim if we added it
         if (input.Shape.Length == 3 && output.Shape.Length == 4 && output.Shape[0] == 1)
@@ -269,11 +264,6 @@ private void ApplyParameters(Vector<T> parameters)
     public override void ResetState()
     {
         _lastInput = null;
-        _bn1Out = null;
-        _relu1Out = null;
-        _conv1Out = null;
-        _bn2Out = null;
-        _relu2Out = null;
 
         _bn1.ResetState();
         _conv1x1.ResetState();

src/NeuralNetworks/Layers/TransitionLayer.cs

@@ -1,4 +1,4 @@
-using AiDotNet.ActivationFunctions;
+using AiDotNet.ActivationFunctions;
 using AiDotNet.Attributes;
 using AiDotNet.Autodiff;
 using AiDotNet.Interfaces;
@@ -82,11 +82,6 @@ public partial class TransitionLayer<T> : LayerBase<T>, ILayerSerializationExtra
     private readonly AveragePoolingLayer<T> _pool;
     private readonly IActivationFunction<T> _relu;
 
-    [Scratch]
-    private Tensor<T>? _lastInput;
-    private Tensor<T>? _bnOut;
-    private Tensor<T>? _reluOut;
-    private Tensor<T>? _convOut;
 
     /// <summary>
     /// Stores the original input shape for any-rank tensor support.
@@ -304,19 +299,18 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
             throw new ArgumentException($"TransitionLayer requires at least 3D input, got {rank}D");
         }
 
-        _lastInput = processInput;
-
-        // BN → ReLU → Conv1x1 → AvgPool
-        _bnOut = _bn.Forward(processInput);
-        _reluOut = _relu.Activate(_bnOut);
-        _convOut = conv.Forward(_reluOut);
+        // BN → ReLU → Conv1x1 → AvgPool. Locals, not fields: nothing reads them after the forward,
+        // and a field would keep this layer's activations alive between forwards.
+        var bnOut = _bn.Forward(processInput);
+        var reluOut = _relu.Activate(bnOut);
+        var convOut = conv.Forward(reluOut);
 
         // Handle batched input (4D) - use Engine.AvgPool2D directly
         // AveragePoolingLayer expects 3D input, so we handle 4D separately
         // [B, C, H, W] uses Engine directly, [C, H, W] uses the AveragePoolingLayer
-        Tensor<T> output = _convOut.Shape.Length == 4
-            ? Engine.AvgPool2D(_convOut, poolSize: 2, stride: 2, padding: 0)
-            : _pool.Forward(_convOut);
+        Tensor<T> output = convOut.Shape.Length == 4
+            ? Engine.AvgPool2D(convOut, poolSize: 2, stride: 2, padding: 0)
+            : _pool.Forward(convOut);
 
         // Restore original batch dimensions for any-rank support
         if (_originalInputShape != null && _originalInputShape.Length > 4)
@@ -429,10 +423,6 @@ internal override Dictionary<string, string> GetMetadata()
     /// </summary>
     public override void ResetState()
     {
-        _lastInput = null;
-        _bnOut = null;
-        _reluOut = null;
-        _convOut = null;
 
         _bn.ResetState();
         _conv?.ResetState();

src/NeuralNetworks/NeuralNetworkBase.cs

@@ -6293,6 +6293,12 @@ private Tensor<T> PredictAndValidateChain(Tensor<T> input)
         {
             observed = InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input);
         }
+        catch
+        {
+            // The recording holds every layer's activations; a failed forward must not leave them alive.
+            trace?.ReleaseRecording();
+            throw;
+        }
         finally
         {
             trace?.Dispose();
@@ -6303,16 +6309,24 @@ private Tensor<T> PredictAndValidateChain(Tensor<T> input)
         // untraced path, and a second forward would draw a stochastic layer's randomness twice (a
         // seeded Monte Carlo dropout call then differed from the next one). Under an enclosing
         // observer the re-run would take the same unfolded path, so it is skipped there too.
-        output = trace is null || !trace.SuppressedFolding || enclosedByObserver
-            ? observed
-            : (InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input));
+        try
+        {
+            output = trace is null || !trace.SuppressedFolding || enclosedByObserver
+                ? observed
+                : (InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input));
 
-        // AFTER the forward, deliberately. Measured with a probe rather than assumed: at every
-        // earlier point tried - ReportLayerContractMismatches, ResolveLazyLayerShapes, and the top
-        // of this method's caller - Layers was still EMPTY, because these stacks materialise during
-        // the forward. A check that runs before the layers exist cannot fail, and would have shipped
-        // as a diagnostic that silently never fired.
-        ReportArchitectureLayerDisagreement();
+            // AFTER the forward, deliberately. Measured with a probe rather than assumed: at every
+            // earlier point tried - ReportLayerContractMismatches, ResolveLazyLayerShapes, and the top
+            // of this method's caller - Layers was still EMPTY, because these stacks materialise during
+            // the forward. A check that runs before the layers exist cannot fail, and would have shipped
+            // as a diagnostic that silently never fired.
+            ReportArchitectureLayerDisagreement();
+        }
+        catch
+        {
+            trace?.ReleaseRecording();
+            throw;
+        }
 
         if (trace is not null)
         {
@@ -6324,6 +6338,13 @@ private Tensor<T> PredictAndValidateChain(Tensor<T> input)
             {
                 // Diagnostic only.
             }
+            finally
+            {
+                // The recording holds every layer's input and output of that forward. Drop it here, so a
+                // context that captured the observer (a thread or task started during the forward) cannot
+                // keep those activations alive.
+                trace.ReleaseRecording();
+            }
         }
 
         return output;