fix: make MCDropoutLayer thread-safe

Author: ooples

src/UncertaintyQuantification/Layers/MCDropoutLayer.cs

@@ -1,5 +1,6 @@
 using AiDotNet.LinearAlgebra;
 using AiDotNet.NeuralNetworks.Layers;
+using System.Threading;
 
 namespace AiDotNet.UncertaintyQuantification.Layers;
 
@@ -27,12 +28,12 @@ namespace AiDotNet.UncertaintyQuantification.Layers;
 /// </remarks>
 public class MCDropoutLayer<T> : LayerBase<T>
 {
-    private Random _rng;
     private readonly double _dropoutRate;
     private readonly T _scale;
     private readonly int? _initialSeed;
-    private Tensor<T>? _lastInput;
-    private Vector<T>? _dropoutMask;
+    private readonly ThreadLocal<Random> _rng;
+    private readonly ThreadLocal<Tensor<T>?> _lastInput = new(() => null);
+    private readonly ThreadLocal<Vector<T>?> _dropoutMask = new(() => null);
     private bool _mcMode; // Monte Carlo mode - always apply dropout
 
     /// <summary>
@@ -69,12 +70,20 @@ public MCDropoutLayer(double dropoutRate = 0.5, bool mcMode = false, int? random
         _dropoutRate = dropoutRate;
         _mcMode = mcMode;
         _initialSeed = randomSeed;
-        _rng = randomSeed.HasValue ? new Random(randomSeed.Value) : new Random();
+        _rng = new ThreadLocal<Random>(() =>
+        {
+            if (_initialSeed.HasValue)
+            {
+                return new Random(unchecked(_initialSeed.Value + Thread.CurrentThread.ManagedThreadId));
+            }
+
+            return new Random();
+        });
     }
 
     internal void ResetRng(int seed)
     {
-        _rng = new Random(seed);
+        _rng.Value = new Random(seed);
     }
 
     /// <summary>
@@ -84,30 +93,32 @@ internal void ResetRng(int seed)
     /// <returns>The output tensor with dropout applied if in training or MC mode.</returns>
     public override Tensor<T> Forward(Tensor<T> input)
     {
-        _lastInput = input;
+        _lastInput.Value = input;
 
         // Apply dropout if in training mode OR Monte Carlo mode
         if (!IsTrainingMode && !_mcMode)
             return input;
 
         var inputVector = input.ToVector();
-        _dropoutMask = new Vector<T>(inputVector.Length);
+        var mask = new Vector<T>(inputVector.Length);
         var outputVector = new Vector<T>(inputVector.Length);
 
         for (int i = 0; i < inputVector.Length; i++)
         {
-            if (_rng.NextDouble() > _dropoutRate)
+            if (_rng.Value!.NextDouble() > _dropoutRate)
             {
-                _dropoutMask[i] = _scale;
+                mask[i] = _scale;
                 outputVector[i] = NumOps.Multiply(inputVector[i], _scale);
             }
             else
             {
-                _dropoutMask[i] = NumOps.Zero;
+                mask[i] = NumOps.Zero;
                 outputVector[i] = NumOps.Zero;
             }
         }
 
+        _dropoutMask.Value = mask;
+
         var outputTensor = Tensor<T>.FromVector(outputVector);
         return input.Shape.Length > 1 ? outputTensor.Reshape(input.Shape) : outputTensor;
     }
@@ -119,7 +130,7 @@ public override Tensor<T> Forward(Tensor<T> input)
     /// <returns>The gradient to pass to the previous layer.</returns>
     public override Tensor<T> Backward(Tensor<T> outputGradient)
     {
-        if (_lastInput == null || _dropoutMask == null)
+        if (_lastInput.Value == null || _dropoutMask.Value == null)
             throw new InvalidOperationException("Forward pass must be called before backward pass.");
 
         if (!IsTrainingMode && !_mcMode)
@@ -130,11 +141,11 @@ public override Tensor<T> Backward(Tensor<T> outputGradient)
 
         for (int i = 0; i < outputGradientVector.Length; i++)
         {
-            inputGradientVector[i] = NumOps.Multiply(outputGradientVector[i], _dropoutMask[i]);
+            inputGradientVector[i] = NumOps.Multiply(outputGradientVector[i], _dropoutMask.Value[i]);
         }
 
         var inputGradientTensor = Tensor<T>.FromVector(inputGradientVector);
-        return _lastInput.Shape.Length > 1 ? inputGradientTensor.Reshape(_lastInput.Shape) : inputGradientTensor;
+        return _lastInput.Value.Shape.Length > 1 ? inputGradientTensor.Reshape(_lastInput.Value.Shape) : inputGradientTensor;
     }
 
     /// <summary>
@@ -169,8 +180,8 @@ public override void SetParameters(Vector<T> parameters)
     /// </summary>
     public override void ResetState()
     {
-        _lastInput = null;
-        _dropoutMask = null;
+        _lastInput.Value = null;
+        _dropoutMask.Value = null;
     }
 
     public override LayerBase<T> Clone()