diff --git a/.gitignore b/.gitignore
index bb704a7df..a9b2d672b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,7 +3,6 @@ test.model
 /vendor
 composer.lock
 .phpunit.result.cache
-.php-cs-fixer.cache
 Thumbs.db
 .DS_Store
 debug.log
diff --git a/composer.json b/composer.json
index 689cb8b0d..0c2137988 100644
--- a/composer.json
+++ b/composer.json
@@ -84,6 +84,9 @@
         "check": [
             "php-cs-fixer fix --config=.php-cs-fixer.dist.php -vvv --dry-run --using-cache=no --sequential --show-progress=dots --stop-on-violation"
         ],
+        "diff": [
+            "php-cs-fixer fix --config=.php-cs-fixer.dist.php -vvv --dry-run --diff"
+        ],
         "fix": [
             "php-cs-fixer fix --config=.php-cs-fixer.dist.php"
         ],
diff --git a/docs/neural-network/activation-functions/softmax.md b/docs/neural-network/activation-functions/softmax.md
index 368ae7ba7..757001c7c 100644
--- a/docs/neural-network/activation-functions/softmax.md
+++ b/docs/neural-network/activation-functions/softmax.md
@@ -1,7 +1,9 @@
-<span style="float:right;"><a href="https://github.com/RubixML/ML/blob/master/src/NeuralNet/ActivationFunctions/Softmax/Softmax.php">[source]</a></span>
+<span style="float:right;"><a href="https://github.com/RubixML/ML/blob/master/src/NeuralNet/ActivationFunctions/Softmax.php">[source]</a></span>
 
 # Softmax
-The Softmax function is a generalization of the [Sigmoid](sigmoid.md) function that squashes each activation between 0 and 1 with the addition that all activations add up to 1. Together, these properties allow the output of the Softmax function to be interpretable as a *joint* probability distribution.
+The Softmax function is a generalization of the [Sigmoid](sigmoid.md) function that squashes each activation between 0 and 1 with the addition that all activations for each sample add up to 1. Together, these properties allow the output of the Softmax function to be interpretable as a *joint* probability distribution for multiclass classification.
+
+Softmax expects batched network activations in `[classes, batch]` layout, where rows represent classes and columns represent samples. Each sample column is normalized independently.
 
 $$
 \text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}
@@ -23,7 +25,7 @@ This activation function does not have any parameters.
 
 ## Example
 ```php
-use Rubix\ML\NeuralNet\ActivationFunctions\Softmax\Softmax;
+use Rubix\ML\NeuralNet\ActivationFunctions\Softmax;
 
 $activationFunction = new Softmax();
 ```
diff --git a/phpunit.xml b/phpunit.xml
index 4680d36cf..f8fbcaeaa 100644
--- a/phpunit.xml
+++ b/phpunit.xml
@@ -3,6 +3,7 @@
   xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
   backupGlobals="false"
   bootstrap="vendor/autoload.php"
+  cacheDirectory="runtime"
   colors="true"
   displayDetailsOnTestsThatTriggerDeprecations="true"
   displayDetailsOnTestsThatTriggerNotices="true"
diff --git a/src/Classifiers/LogisticRegression.php b/src/Classifiers/LogisticRegression.php
index 67329138e..f3a7d682a 100644
--- a/src/Classifiers/LogisticRegression.php
+++ b/src/Classifiers/LogisticRegression.php
@@ -2,6 +2,8 @@
 
 namespace Rubix\ML\Classifiers;
 
+use Generator;
+use NumPower;
 use Rubix\ML\Online;
 use Rubix\ML\Learner;
 use Rubix\ML\Verbose;
@@ -33,7 +35,6 @@
 use Rubix\ML\Specifications\SamplesAreCompatibleWithEstimator;
 use Rubix\ML\Exceptions\InvalidArgumentException;
 use Rubix\ML\Exceptions\RuntimeException;
-use Generator;
 
 use function is_nan;
 use function count;
@@ -429,7 +430,7 @@ public function proba(Dataset $dataset) : array
 
         $activations = $this->network->infer($dataset);
 
-        $activations = array_column($activations->asArray(), 0);
+        $activations = array_column($activations->toArray(), 0);
 
         $probabilities = [];
 
@@ -461,10 +462,9 @@ public function featureImportances() : array
             throw new RuntimeException('Weight layer not found.');
         }
 
-        return $layer->weights()
-            ->rowAsVector(0)
-            ->abs()
-            ->asArray();
+        $weights = NumPower::abs($layer->weights())->toArray();
+
+        return $weights[0] ?? [];
     }
 
     /**
diff --git a/src/Classifiers/MultilayerPerceptron.php b/src/Classifiers/MultilayerPerceptron.php
index e62f6e206..586be1dc6 100644
--- a/src/Classifiers/MultilayerPerceptron.php
+++ b/src/Classifiers/MultilayerPerceptron.php
@@ -545,7 +545,7 @@ public function proba(Dataset $dataset) : array
 
         $probabilities = [];
 
-        foreach ($activations->asArray() as $dist) {
+        foreach ($activations->toArray() as $dist) {
             $probabilities[] = array_combine($this->classes, $dist) ?: [];
         }
 
diff --git a/src/Classifiers/SoftmaxClassifier.php b/src/Classifiers/SoftmaxClassifier.php
index 99f227564..ca4f782fc 100644
--- a/src/Classifiers/SoftmaxClassifier.php
+++ b/src/Classifiers/SoftmaxClassifier.php
@@ -424,7 +424,7 @@ public function proba(Dataset $dataset) : array
 
         $probabilities = [];
 
-        foreach ($activations->asArray() as $dist) {
+        foreach ($activations->toArray() as $dist) {
             $probabilities[] = array_combine($this->classes, $dist) ?: [];
         }
 
diff --git a/src/NeuralNet/ActivationFunctions/Softmax.php b/src/NeuralNet/ActivationFunctions/Softmax.php
index 090dd402f..5b9dafdc4 100644
--- a/src/NeuralNet/ActivationFunctions/Softmax.php
+++ b/src/NeuralNet/ActivationFunctions/Softmax.php
@@ -13,6 +13,8 @@
  * The Softmax function is a generalization of the Sigmoid function that squashes
  * each activation between 0 and 1, and all activations add up to 1.
  *
+ * Expects network layout `[classes, batch]` and normalizes each sample column.
+ *
  * @category    Machine Learning
  * @package     Rubix/ML
  * @author      Andrew DalPino
@@ -26,37 +28,34 @@ class Softmax implements ActivationFunction, OBufferDerivative
      * The Softmax function is defined as:
      * f(x_i) = exp(x_i) / sum(exp(x_j)) for all j
      *
-     * The Softmax function is a generalization of the Sigmoid function that squashes
-     * each activation between 0 and 1, and all activations add up to 1.
-     *
-     * > **Note:** This function can be rewritten in a more efficient way,
-     * using NumPower::exp(), NumPower::sum(), and NumPower::divide().
-     * Currently blocked by implementation of 2nd parameter "axis" for NumPower::sum()
+     * Numerically stable form subtracts the per-sample max before exponentiation.
      *
      * @param NDArray $input
      * @return NDArray
      */
     public function activate(NDArray $input) : NDArray
     {
-        // Convert to PHP array for stable processing
-        $inputArray = $input->toArray();
-        $result = [];
+        $columns = $input->shape()[1];
+        $values = $input->toArray();
 
-        // Process each row separately to ensure row-wise normalization
-        foreach ($inputArray as $row) {
-            $expRow = array_map('exp', $row);
-            $sum = array_sum($expRow);
-            $softmaxRow = [];
+        // NumPower::max() has no axis argument, so compute column maxima in PHP.
+        $maxima = [];
 
-            foreach ($expRow as $value) {
-                // Round to 7 decimal places to match test expectations
-                $softmaxRow[] = round($value / $sum, 7);
+        for ($column = 0; $column < $columns; ++$column) {
+            $maximum = -INF;
+
+            foreach ($values as $row) {
+                $maximum = max($maximum, $row[$column]);
             }
 
-            $result[] = $softmaxRow;
+            $maxima[] = $maximum;
         }
 
-        return NumPower::array($result);
+        $max = NumPower::reshape(NumPower::array($maxima), [1, $columns]);
+        $exponentials = NumPower::exp(NumPower::subtract($input, $max));
+        $totals = NumPower::reshape(NumPower::sum($exponentials, axis: 0), [1, $columns]);
+
+        return NumPower::divide($exponentials, $totals);
     }
 
     /**
diff --git a/src/NeuralNet/FeedForward.php b/src/NeuralNet/FeedForward.php
index caaf6890f..01efc301b 100644
--- a/src/NeuralNet/FeedForward.php
+++ b/src/NeuralNet/FeedForward.php
@@ -279,4 +279,3 @@ public function exportGraphviz() : Encoding
         return new Encoding($dot);
     }
 }
-
diff --git a/src/NeuralNet/Layers/Multiclass.php b/src/NeuralNet/Layers/Multiclass.php
index 6e238b967..ec03e148e 100644
--- a/src/NeuralNet/Layers/Multiclass.php
+++ b/src/NeuralNet/Layers/Multiclass.php
@@ -158,16 +158,17 @@ public function back(array $labels, Optimizer $optimizer) : array
                 . ' before backpropagating.');
         }
 
+        // Build one-hot targets as [classes, batch] to match Dense output layout.
         $expected = [];
 
-        foreach ($labels as $label) {
-            $dist = [];
+        foreach ($this->classes as $class) {
+            $row = [];
 
-            foreach ($this->classes as $class) {
-                $dist[] = $class == $label ? 1.0 : 0.0;
+            foreach ($labels as $label) {
+                $row[] = $class == $label ? 1.0 : 0.0;
             }
 
-            $expected[] = $dist;
+            $expected[] = $row;
         }
 
         $expected = NumPower::array($expected);
diff --git a/tests/Classifiers/LogisticRegressionTest.php b/tests/Classifiers/LogisticRegressionTest.php
index 27e7fa87c..292bf6076 100644
--- a/tests/Classifiers/LogisticRegressionTest.php
+++ b/tests/Classifiers/LogisticRegressionTest.php
@@ -162,7 +162,7 @@ public function testTrainPartialPredict() : void
 
         $this->assertGreaterThanOrEqual(self::MIN_SCORE, $score);
 
-        $this->assertEquals('58a6bb3c', $this->estimator->revision());
+        $this->assertEquals('f2e08c1a', $this->estimator->revision());
     }
 
     public function testTrainIncompatible() : void
diff --git a/tests/NeuralNet/ActivationFunctions/SoftmaxTest.php b/tests/NeuralNet/ActivationFunctions/SoftmaxTest.php
index 5d02b0cae..9c13865c0 100644
--- a/tests/NeuralNet/ActivationFunctions/SoftmaxTest.php
+++ b/tests/NeuralNet/ActivationFunctions/SoftmaxTest.php
@@ -14,7 +14,6 @@
 use PHPUnit\Framework\Attributes\TestDox;
 use PHPUnit\Framework\TestCase;
 use Rubix\ML\NeuralNet\ActivationFunctions\Softmax;
-use Tensor\Matrix;
 
 #[Group('ActivationFunctions')]
 #[CoversClass(Softmax::class)]
@@ -30,46 +29,58 @@ class SoftmaxTest extends TestCase
      */
     public static function computeProvider() : Generator
     {
+        // Inputs use network layout [classes, batch].
         yield [
             NumPower::array([
-                [2.0, 1.0, -0.5, 0.0],
+                [2.0],
+                [1.0],
+                [-0.5],
+                [0.0],
             ]),
             [
-                [0.6307954, 0.2320567, 0.0517789, 0.0853689],
+                [0.6307955],
+                [0.2320567],
+                [0.0517788],
+                [0.0853688],
             ],
         ];
 
         yield [
             NumPower::array([
-                [-0.12, 0.31, -0.49],
-                [0.99, 0.08, -0.03],
-                [0.05, -0.52, 0.54],
+                [-0.12, 0.99, 0.05],
+                [0.31, 0.08, -0.52],
+                [-0.49, -0.03, 0.54],
             ]),
             [
-                [0.3097901, 0.4762271, 0.2139827],
-                [0.5671765, 0.2283022, 0.2045210],
-                [0.312711, 0.176846, 0.510443],
+                [0.3097901, 0.5671766, 0.3127109],
+                [0.4762272, 0.2283023, 0.1768459],
+                [0.2139826, 0.2045210, 0.5104430],
             ],
         ];
 
-        // Test with zeros
         yield [
             NumPower::array([
-                [0.0, 0.0, 0.0, 0.0],
+                [0.0],
+                [0.0],
+                [0.0],
+                [0.0],
             ]),
             [
-                [0.25, 0.25, 0.25, 0.25],
+                [0.25],
+                [0.25],
+                [0.25],
+                [0.25],
             ],
         ];
 
         yield [
             NumPower::array([
-                [1, 2],
-                [3, 4],
+                [1, 3],
+                [2, 4],
             ]),
             [
-                [0.2689414, 0.7310585],
-                [0.2689414, 0.7310585],
+                [0.2689414, 0.2689414],
+                [0.7310585, 0.7310585],
             ],
         ];
     }
@@ -79,10 +90,10 @@ public static function computeProvider() : Generator
      */
     public static function differentiateProvider() : Generator
     {
-        // Test with simple values
         yield [
             NumPower::array([
-                [0.6, 0.4],
+                [0.6],
+                [0.4],
             ]),
             [
                 [0.24, -0.24],
@@ -90,10 +101,11 @@ public static function differentiateProvider() : Generator
             ],
         ];
 
-        // Test with more complex values
         yield [
             NumPower::array([
-                [0.3, 0.5, 0.2],
+                [0.3],
+                [0.5],
+                [0.2],
             ]),
             [
                 [0.21, -0.15, -0.06],
@@ -102,10 +114,10 @@ public static function differentiateProvider() : Generator
             ],
         ];
 
-        // Test 2x2 matrix
         yield [
             NumPower::array([
-                [0.2689414, 0.7310585],
+                [0.2689414],
+                [0.7310585],
             ]),
             [
                 [0.1966119, -0.19661192],
@@ -119,24 +131,29 @@ public static function differentiateProvider() : Generator
      */
     public static function sumToOneProvider() : Generator
     {
-        // Test with various input values
         yield [
             NumPower::array([
-                [10.0, -5.0, 3.0, 2.0],
+                [10.0],
+                [-5.0],
+                [3.0],
+                [2.0],
             ]),
         ];
 
         yield [
             NumPower::array([
-                [-10.0, -20.0, -30.0],
+                [-10.0],
+                [-20.0],
+                [-30.0],
             ]),
         ];
 
         yield [
             NumPower::array([
-                [0.1, 0.2, 0.3, 0.4],
-                [5.0, 4.0, 3.0, 2.0],
-                [-1.0, -2.0, -3.0, -4.0],
+                [0.1, 5.0, -1.0],
+                [0.2, 4.0, -2.0],
+                [0.3, 3.0, -3.0],
+                [0.4, 2.0, -4.0],
             ]),
         ];
     }
@@ -183,15 +200,17 @@ public function testDifferentiate(NDArray $output, array $expected) : void
     #[DataProvider('sumToOneProvider')]
     public function testSumToOne(NDArray $input) : void
     {
-        $activations = $this->activationFn->activate($input);
+        $activations = $this->activationFn->activate($input)->toArray();
+
+        $columns = count($activations[0]);
 
-        // Convert to array for easier processing
-        $activationsArray = $activations->toArray();
+        for ($column = 0; $column < $columns; ++$column) {
+            $sum = 0.0;
+
+            foreach ($activations as $row) {
+                $sum += $row[$column];
+            }
 
-        // Check that each row sums to 1
-        foreach ($activationsArray as $row) {
-            $sum = array_sum($row);
-            // Use a slightly larger delta to account for rounding errors
             static::assertEqualsWithDelta(1.0, $sum, 1e-7);
         }
     }
diff --git a/tests/NeuralNet/Initializers/LeCunNormalTest.php b/tests/NeuralNet/Initializers/LeCunNormalTest.php
index 9f9ce2a39..4c409503c 100644
--- a/tests/NeuralNet/Initializers/LeCunNormalTest.php
+++ b/tests/NeuralNet/Initializers/LeCunNormalTest.php
@@ -140,7 +140,7 @@ public function testDistributionStatisticsMatchLeCunNormal(int $fanIn, int $fanO
         $this->assertThat(
             $std,
             $this->logicalAnd(
-                $this->greaterThan($expectedStd * 0.85),
+                $this->greaterThan($expectedStd * 0.80),
                 $this->lessThan($expectedStd * 1.1)
             ),
             'Standard deviation does not match Le Cun initialization'
diff --git a/tests/NeuralNet/Initializers/TruncatedNormalTest.php b/tests/NeuralNet/Initializers/TruncatedNormalTest.php
index 7a6032916..7cf3ec597 100644
--- a/tests/NeuralNet/Initializers/TruncatedNormalTest.php
+++ b/tests/NeuralNet/Initializers/TruncatedNormalTest.php
@@ -172,7 +172,7 @@ public function testValuesFollowTruncatedNormalDistribution(int $fanIn, int $fan
         $this->assertThat(
             $resultStd,
             $this->logicalAnd(
-                $this->greaterThan($stdDev * 0.85),
+                $this->greaterThan($stdDev * 0.80),
                 $this->lessThan($stdDev * 1.1)
             ),
             'Standard deviation does not match Truncated Normal initialization'
diff --git a/tests/NeuralNet/Layers/MulticlassTest.php b/tests/NeuralNet/Layers/MulticlassTest.php
index be56d9442..a0ab7b89c 100644
--- a/tests/NeuralNet/Layers/MulticlassTest.php
+++ b/tests/NeuralNet/Layers/MulticlassTest.php
@@ -51,9 +51,9 @@ public static function forwardProvider() : array
     {
         return [
             'expectedForward' => [[
-                [0.1719820, 0.7707700, 0.0572478],
-                [0.0498033, 0.0450639, 0.9051327],
-                [0.6219707, 0.0015385, 0.3764905],
+                [0.1719820, 0.0498033, 0.6219707],
+                [0.7707700, 0.0450639, 0.0015386],
+                [0.0572478, 0.9051328, 0.3764906],
             ]],
         ];
     }
@@ -65,9 +65,9 @@ public static function backProvider() : array
     {
         return [
             'expectedGradient' => [[
-                [-0.0920019, 0.0856411, 0.0063608],
-                [0.0055337, -0.1061040, 0.1005703],
-                [0.0691078, 0.00017093, -0.0692788],
+                [-0.0920019, 0.0055337, 0.0691078],
+                [0.0856411, -0.1061040, 0.0001709],
+                [0.0063608, 0.1005703, -0.0692788],
             ]],
         ];
     }
@@ -83,10 +83,11 @@ public static function inferProvider() : array
 
     protected function setUp() : void
     {
+        // Column layout [classes, batch] matching Dense / FeedForward.
         $this->input = NumPower::array([
-            [1.0, 2.5, -0.1],
-            [0.1, 0.0, 3.0],
-            [0.002, -6.0, -0.5],
+            [1.0, 0.1, 0.002],
+            [2.5, 0.0, -6.0],
+            [-0.1, 3.0, -0.5],
         ]);
 
         $this->labels = ['hot', 'cold', 'ice cold'];
@@ -177,14 +178,14 @@ public function testGradient(array $expectedGradient) : void
         // Rebuild expected one-hot matrix the same way as Multiclass::back()
         $expected = [];
 
-        foreach ($this->labels as $label) {
-            $dist = [];
+        foreach (['hot', 'cold', 'ice cold'] as $class) {
+            $row = [];
 
-            foreach (['hot', 'cold', 'ice cold'] as $class) {
-                $dist[] = $class === $label ? 1.0 : 0.0;
+            foreach ($this->labels as $label) {
+                $row[] = $class === $label ? 1.0 : 0.0;
             }
 
-            $expected[] = $dist;
+            $expected[] = $row;
         }
 
         $expectedNd = NumPower::array($expected);
diff --git a/tests/Regressors/RidgeTest.php b/tests/Regressors/RidgeTest.php
index a5ddbf832..fdaa6a2cd 100644
--- a/tests/Regressors/RidgeTest.php
+++ b/tests/Regressors/RidgeTest.php
@@ -10,8 +10,6 @@
 use PHPUnit\Framework\Attributes\Test;
 use PHPUnit\Framework\Attributes\TestDox;
 use PHPUnit\Framework\TestCase;
-use NumPower;
-use ReflectionClass;
 use Rubix\ML\CrossValidation\Metrics\RSquared;
 use Rubix\ML\Datasets\Generators\Hyperplane;
 use Rubix\ML\Datasets\Labeled;
@@ -466,6 +464,11 @@ public function randomDatasetsProduceFinitePredictions() : void
     }
 
     /**
+     * Make random linear problem
+     *
+     * @param int $samples
+     * @param int $features
+     * @param int $seed
      * @return array{0: list<list<float>>, 1: list<float>}
      */
     private function makeRandomLinearProblem(int $samples, int $features, int $seed) : array