fix: address PR review comments for pruning strategies

- PruningMask: use PointwiseMultiply vectorization for SIMD acceleration
- StructuredPruningStrategy: implement Filter and Channel pruning types
- LotteryTicketPruningStrategy: O(n⁴)→O(n²) optimization, add validation
- PruningStrategyTests: fix potential integer overflow
- ModelCompressionBenchmarks: fix weight-matrix sizing losing tail elements

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: franklinic

AiDotNetBenchmarkTests/BenchmarkTests/ModelCompressionBenchmarks.cs

@@ -61,13 +61,19 @@ public void Setup()
         }
 
         // Initialize weight matrix (for structured pruning)
+        // Use ceiling division to ensure all weights fit in the matrix
         int rows = (int)Math.Sqrt(WeightCount);
-        int cols = WeightCount / rows;
+        int cols = (WeightCount + rows - 1) / rows; // Ceiling division: ensures rows * cols >= WeightCount
         _weightMatrix = new Matrix<double>(rows, cols);
-        int idx = 0;
-        for (int i = 0; i < rows && idx < WeightCount; i++)
-            for (int j = 0; j < cols && idx < WeightCount; j++)
-                _weightMatrix[i, j] = _weights[idx++];
+        for (int i = 0; i < rows; i++)
+        {
+            for (int j = 0; j < cols; j++)
+            {
+                int idx = i * cols + j;
+                // Fill with weight data if available, otherwise zero (padding for alignment)
+                _weightMatrix[i, j] = idx < WeightCount ? _weights[idx] : 0.0;
+            }
+        }
 
         // Initialize compression algorithms
         _deepCompression = new DeepCompression<double>(

src/Pruning/LotteryTicketPruningStrategy.cs

@@ -100,6 +100,9 @@ public class LotteryTicketPruningStrategy<T> : IPruningStrategy<T>
     /// </remarks>
     public LotteryTicketPruningStrategy(int iterativeRounds = 5)
     {
+        if (iterativeRounds <= 0)
+            throw new ArgumentOutOfRangeException(nameof(iterativeRounds), "iterativeRounds must be greater than 0.");
+
         _numOps = MathHelper.GetNumericOperations<T>();
         _initialWeights = new Dictionary<string, Matrix<T>>();
         _iterativeRounds = iterativeRounds;
@@ -134,10 +137,10 @@ public void StoreInitialWeights(string layerName, Matrix<T> weights)
     /// <exception cref="InvalidOperationException">Thrown when no weights stored for the layer</exception>
     public Matrix<T> GetInitialWeights(string layerName)
     {
-        if (!_initialWeights.ContainsKey(layerName))
+        if (!_initialWeights.TryGetValue(layerName, out var weights))
             throw new InvalidOperationException($"No initial weights stored for layer {layerName}");
 
-        return _initialWeights[layerName].Clone();
+        return weights.Clone();
     }
 
     /// <summary>
@@ -210,6 +213,9 @@ public Tensor<T> ComputeImportanceScores(Tensor<T> weights, Tensor<T>? gradients
     /// <returns>Binary mask (1 = keep, 0 = prune)</returns>
     public IPruningMask<T> CreateMask(Vector<T> importanceScores, double targetSparsity)
     {
+        if (targetSparsity < 0.0 || targetSparsity > 1.0)
+            throw new ArgumentException("targetSparsity must be between 0 and 1 (inclusive).", nameof(targetSparsity));
+
         double prunePerRound = 1.0 - Math.Pow(1.0 - targetSparsity, 1.0 / _iterativeRounds);
         var currentMask = new PruningMask<T>(1, importanceScores.Length);
 
@@ -230,9 +236,10 @@ public IPruningMask<T> CreateMask(Vector<T> importanceScores, double targetSpars
 
             flatScores.Sort((a, b) => a.score.CompareTo(b.score));
 
+            // Reuse maskedScores instead of recomputing currentMask.Apply(importanceScores) in loop
             var keepIndices = new bool[importanceScores.Length];
             for (int i = 0; i < importanceScores.Length; i++)
-                keepIndices[i] = !_numOps.Equals(currentMask.Apply(importanceScores)[i], _numOps.Zero);
+                keepIndices[i] = !_numOps.Equals(maskedScores[i], _numOps.Zero);
 
             for (int i = 0; i < numToPrune && i < flatScores.Count; i++)
             {
@@ -267,6 +274,9 @@ public IPruningMask<T> CreateMask(Vector<T> importanceScores, double targetSpars
     /// </remarks>
     public IPruningMask<T> CreateMask(Matrix<T> importanceScores, double targetSparsity)
     {
+        if (targetSparsity < 0.0 || targetSparsity > 1.0)
+            throw new ArgumentException("targetSparsity must be between 0 and 1 (inclusive).", nameof(targetSparsity));
+
         // Iterative magnitude pruning to target sparsity
         // Each round prunes (1 - (1 - targetSparsity)^(1/rounds)) of remaining weights
         double prunePerRound = 1.0 - Math.Pow(1.0 - targetSparsity, 1.0 / _iterativeRounds);
@@ -298,11 +308,12 @@ public IPruningMask<T> CreateMask(Matrix<T> importanceScores, double targetSpars
 
             flatScores.Sort((a, b) => a.score.CompareTo(b.score));
 
+            // Reuse maskedScores instead of recomputing currentMask.Apply(importanceScores) in loop (O(n⁴) → O(n²))
             var keepIndices = new bool[importanceScores.Rows, importanceScores.Columns];
 
             for (int i = 0; i < importanceScores.Rows; i++)
                 for (int j = 0; j < importanceScores.Columns; j++)
-                    keepIndices[i, j] = !_numOps.Equals(currentMask.Apply(importanceScores)[i, j], _numOps.Zero);
+                    keepIndices[i, j] = !_numOps.Equals(maskedScores[i, j], _numOps.Zero);
 
             for (int i = 0; i < numToPrune && i < flatScores.Count; i++)
             {
@@ -324,6 +335,9 @@ public IPruningMask<T> CreateMask(Matrix<T> importanceScores, double targetSpars
     /// <returns>Binary mask (1 = keep, 0 = prune)</returns>
     public IPruningMask<T> CreateMask(Tensor<T> importanceScores, double targetSparsity)
     {
+        if (targetSparsity < 0.0 || targetSparsity > 1.0)
+            throw new ArgumentException("targetSparsity must be between 0 and 1 (inclusive).", nameof(targetSparsity));
+
         double prunePerRound = 1.0 - Math.Pow(1.0 - targetSparsity, 1.0 / _iterativeRounds);
         var flatScoresInit = importanceScores.ToVector();
         // Initialize mask to all true
@@ -616,14 +630,16 @@ public void ResetToInitialWeights(string layerName, Matrix<T> weights, IPruningM
         if (initial.Rows != weights.Rows || initial.Columns != weights.Columns)
             throw new ArgumentException("Weight dimensions don't match initial weights");
 
+        // Compute masked initial weights once (O(n²) instead of O(n⁴))
+        var maskedInitial = mask.Apply(initial);
+
         // Reset non-pruned weights to their initialization
         for (int i = 0; i < weights.Rows; i++)
         {
             for (int j = 0; j < weights.Columns; j++)
             {
                 // Keep initial value where mask is 1, zero otherwise
-                var maskValue = mask.Apply(initial)[i, j];
-                weights[i, j] = maskValue;
+                weights[i, j] = maskedInitial[i, j];
             }
         }
     }

src/Pruning/PruningMask.cs

@@ -1,3 +1,4 @@
+using AiDotNet.Extensions;
 using AiDotNet.Interfaces;
 
 namespace AiDotNet.Pruning;
@@ -158,17 +159,8 @@ public Matrix<T> Apply(Matrix<T> weights)
         if (weights.Rows != _mask.Rows || weights.Columns != _mask.Columns)
             throw new ArgumentException("Weight matrix shape must match mask shape");
 
-        var result = new Matrix<T>(weights.Rows, weights.Columns);
-
-        for (int i = 0; i < weights.Rows; i++)
-        {
-            for (int j = 0; j < weights.Columns; j++)
-            {
-                result[i, j] = _numOps.Multiply(weights[i, j], _mask[i, j]);
-            }
-        }
-
-        return result;
+        // Use vectorized PointwiseMultiply for SIMD acceleration
+        return weights.PointwiseMultiply(_mask);
     }
 
     /// <summary>
@@ -195,29 +187,69 @@ public Tensor<T> Apply(Tensor<T> weights)
         // For 4D tensors (convolutional layers: [filters, channels, height, width])
         if (weights.Rank == 4)
         {
-            // Clone the tensor manually since Clone() might not exist
-            var flatData = weights.ToVector();
-            var result = Tensor<T>.FromVector(flatData, (int[])weights.Shape.Clone());
             int filters = weights.Shape[0];
             int channels = weights.Shape[1];
+            int height = weights.Shape[2];
+            int width = weights.Shape[3];
+            int spatialSize = height * width;
 
-            // Apply mask element-wise for now (unstructured pruning)
-            // For structured pruning, this would need to be modified
-            for (int f = 0; f < filters; f++)
+            // Check if mask dimensions match filter/channel for structured pruning
+            if (_mask.Rows == filters && _mask.Columns == channels)
             {
-                for (int c = 0; c < channels; c++)
+                // Apply filter/channel-level pruning using vectorized operations
+                // Broadcast mask [filters, channels] to [filters, channels, height, width]
+                int totalElements = filters * channels * spatialSize;
+                var maskData = new T[totalElements];
+
+                // Fill mask data by broadcasting each mask value across spatial dimensions
+                int idx = 0;
+                for (int f = 0; f < filters; f++)
                 {
-                    for (int h = 0; h < weights.Shape[2]; h++)
+                    for (int c = 0; c < channels; c++)
                     {
-                        for (int w = 0; w < weights.Shape[3]; w++)
+                        T maskValue = _mask[f, c];
+                        // Fill spatial elements with the same mask value
+                        for (int s = 0; s < spatialSize; s++)
                         {
-                            result[f, c, h, w] = weights[f, c, h, w];
+                            maskData[idx++] = maskValue;
                         }
                     }
                 }
+
+                var maskTensor = Tensor<T>.FromVector(new Vector<T>(maskData), new int[] { filters, channels, height, width });
+
+                // Use vectorized PointwiseMultiply for SIMD acceleration
+                return weights.PointwiseMultiply(maskTensor);
             }
+            else
+            {
+                // Unstructured pruning: apply mask element-by-element
+                var flatWeights = weights.ToVector();
+                int totalElements = flatWeights.Length;
+
+                if (_mask.Rows * _mask.Columns != totalElements)
+                {
+                    throw new ArgumentException(
+                        $"Mask shape [{_mask.Rows}, {_mask.Columns}] does not match 4D tensor total elements ({totalElements}) " +
+                        $"or filter/channel dimensions [{filters}, {channels}]");
+                }
 
-            return result;
+                // Convert mask to flat vector and use vectorized PointwiseMultiply
+                var flatMask = new T[totalElements];
+                int idx = 0;
+                for (int i = 0; i < _mask.Rows; i++)
+                {
+                    for (int j = 0; j < _mask.Columns; j++)
+                    {
+                        flatMask[idx++] = _mask[i, j];
+                    }
+                }
+
+                var flatMaskVector = new Vector<T>(flatMask);
+                var flatResult = flatWeights.PointwiseMultiply(flatMaskVector);
+
+                return Tensor<T>.FromVector(flatResult, (int[])weights.Shape.Clone());
+            }
         }
 
         throw new NotSupportedException($"Tensor rank {weights.Rank} not supported for pruning");
@@ -237,21 +269,12 @@ public Tensor<T> Apply(Tensor<T> weights)
     public Vector<T> Apply(Vector<T> weights)
     {
         // For 1D vectors, the mask is stored as a single-row matrix
-        int length = Math.Min(weights.Length, _mask.Columns);
-        var result = new T[weights.Length];
-
-        for (int i = 0; i < length; i++)
-        {
-            result[i] = _numOps.Multiply(weights[i], _mask[0, i]);
-        }
-
-        // Keep remaining elements as-is if vector is longer than mask
-        for (int i = length; i < weights.Length; i++)
-        {
-            result[i] = weights[i];
-        }
+        if (weights.Length != _mask.Columns)
+            throw new ArgumentException($"Weight vector length ({weights.Length}) must match mask columns ({_mask.Columns})");
 
-        return new Vector<T>(result);
+        // Extract mask row as vector and use vectorized PointwiseMultiply
+        var maskVector = _mask.GetRow(0);
+        return weights.PointwiseMultiply(maskVector);
     }
 
     /// <summary>

src/Pruning/StructuredPruningStrategy.cs

@@ -194,8 +194,47 @@ public Matrix<T> ComputeImportanceScores(Matrix<T> weights, Matrix<T>? gradients
                 }
                 break;
 
-            default:
-                throw new NotImplementedException($"Pruning type {_pruningType} not yet implemented");
+            case StructurePruningType.Filter:
+                // For 2D matrix, Filter pruning treats rows as filters
+                // Score for each filter (row) = L2 norm of its weights
+                for (int row = 0; row < weights.Rows; row++)
+                {
+                    double rowNorm = 0;
+                    for (int col = 0; col < weights.Columns; col++)
+                    {
+                        double val = Convert.ToDouble(weights[row, col]);
+                        rowNorm += val * val;
+                    }
+                    rowNorm = Math.Sqrt(rowNorm);
+
+                    // Assign same score to all weights in row
+                    for (int col = 0; col < weights.Columns; col++)
+                    {
+                        scores[row, col] = _numOps.FromDouble(rowNorm);
+                    }
+                }
+                break;
+
+            case StructurePruningType.Channel:
+                // For 2D matrix, Channel pruning treats columns as channels (same as Neuron)
+                // Score for each channel (column) = L2 norm of its weights
+                for (int col = 0; col < weights.Columns; col++)
+                {
+                    double columnNorm = 0;
+                    for (int row = 0; row < weights.Rows; row++)
+                    {
+                        double val = Convert.ToDouble(weights[row, col]);
+                        columnNorm += val * val;
+                    }
+                    columnNorm = Math.Sqrt(columnNorm);
+
+                    // Assign same score to all weights in column
+                    for (int row = 0; row < weights.Rows; row++)
+                    {
+                        scores[row, col] = _numOps.FromDouble(columnNorm);
+                    }
+                }
+                break;
         }
 
         return scores;
@@ -230,29 +269,30 @@ public IPruningMask<T> CreateMask(Matrix<T> importanceScores, double targetSpars
         switch (_pruningType)
         {
             case StructurePruningType.Neuron:
-                // Prune entire columns (neurons)
-                int totalNeurons = importanceScores.Columns;
-                int neuronsToPrune = (int)(totalNeurons * targetSparsity);
+            case StructurePruningType.Channel:
+                // Prune entire columns (neurons/channels)
+                int totalColumns = importanceScores.Columns;
+                int columnsToPrune = (int)(totalColumns * targetSparsity);
 
-                // Get score for each neuron (all rows in column have same score)
-                var neuronScores = new List<(int col, double score)>();
+                // Get score for each column (all rows in column have same score)
+                var columnScores = new List<(int col, double score)>();
                 for (int col = 0; col < importanceScores.Columns; col++)
                 {
                     double score = Convert.ToDouble(importanceScores[0, col]);
-                    neuronScores.Add((col, score));
+                    columnScores.Add((col, score));
                 }
 
-                // Sort by score (ascending)
-                neuronScores.Sort((a, b) => a.score.CompareTo(b.score));
+                // Sort by score (ascending - lowest scores get pruned first)
+                columnScores.Sort((a, b) => a.score.CompareTo(b.score));
 
                 // Mark columns to keep
                 var keepColumns = new bool[importanceScores.Columns];
                 for (int i = 0; i < importanceScores.Columns; i++)
                     keepColumns[i] = true;
 
-                for (int i = 0; i < neuronsToPrune; i++)
+                for (int i = 0; i < columnsToPrune; i++)
                 {
-                    keepColumns[neuronScores[i].col] = false;
+                    keepColumns[columnScores[i].col] = false;
                 }
 
                 // Create mask
@@ -265,8 +305,41 @@ public IPruningMask<T> CreateMask(Matrix<T> importanceScores, double targetSpars
                 }
                 break;
 
-            default:
-                throw new NotImplementedException($"Pruning type {_pruningType} not yet implemented");
+            case StructurePruningType.Filter:
+                // Prune entire rows (filters)
+                int totalRows = importanceScores.Rows;
+                int rowsToPrune = (int)(totalRows * targetSparsity);
+
+                // Get score for each row (all columns in row have same score)
+                var rowScores = new List<(int row, double score)>();
+                for (int row = 0; row < importanceScores.Rows; row++)
+                {
+                    double score = Convert.ToDouble(importanceScores[row, 0]);
+                    rowScores.Add((row, score));
+                }
+
+                // Sort by score (ascending - lowest scores get pruned first)
+                rowScores.Sort((a, b) => a.score.CompareTo(b.score));
+
+                // Mark rows to keep
+                var keepRows = new bool[importanceScores.Rows];
+                for (int i = 0; i < importanceScores.Rows; i++)
+                    keepRows[i] = true;
+
+                for (int i = 0; i < rowsToPrune; i++)
+                {
+                    keepRows[rowScores[i].row] = false;
+                }
+
+                // Create mask
+                for (int row = 0; row < importanceScores.Rows; row++)
+                {
+                    for (int col = 0; col < importanceScores.Columns; col++)
+                    {
+                        keepIndices[row, col] = keepRows[row];
+                    }
+                }
+                break;
         }
 
         var mask = new PruningMask<T>(importanceScores.Rows, importanceScores.Columns);

tests/AiDotNet.Tests/Pruning/PruningStrategyTests.cs

@@ -343,7 +343,7 @@ public void LotteryTicket_IterativePruning_AchievesTargetSparsity()
         var weights = new Matrix<double>(5, 5);
         for (int i = 0; i < 5; i++)
             for (int j = 0; j < 5; j++)
-                weights[i, j] = (i + 1) * (j + 1) * 0.1; // Varying magnitudes
+                weights[i, j] = ((double)(i + 1)) * (j + 1) * 0.1; // Varying magnitudes
 
         var strategy = new LotteryTicketPruningStrategy<double>(iterativeRounds: 5);