Bug fix for Tensor.logSumExp. (#551)

Fixed a bug in the 'logsumexp' implementation and augment the tests to ensure correctness.

Author: eaplatanios

Sources/TensorFlow/Operators/Math.swift

@@ -2371,8 +2371,8 @@ public extension Tensor where Scalar: TensorFlowFloatingPoint {
     func logSumExp(squeezingAxes axes: Tensor<Int32>) -> Tensor {
         let rawMax = max(alongAxes: axes)
         let offset = withoutDerivative(at: rawMax) { rawMax in 
-            rawMax.replacing(
-                with: Tensor<Scalar>(zerosLike: rawMax),
+            Tensor<Scalar>(zerosLike: rawMax).replacing(
+                with: rawMax,
                 where: rawMax.isFinite)
         }
         let result = TensorFlow.log(TensorFlow.exp(self - offset).sum(squeezingAxes: axes))
@@ -2435,8 +2435,8 @@ public extension Tensor where Scalar: TensorFlowFloatingPoint {
     func logSumExp(alongAxes axes: Tensor<Int32>) -> Tensor {
         let rawMax = max(alongAxes: axes)
         let offset = withoutDerivative(at: rawMax) { rawMax in 
-            rawMax.replacing(
-                with: Tensor<Scalar>(zerosLike: rawMax),
+            Tensor<Scalar>(zerosLike: rawMax).replacing(
+                with: rawMax,
                 where: rawMax.isFinite)
         }
         let result = TensorFlow.log(TensorFlow.exp(self - offset).sum(alongAxes: axes))

Tests/TensorFlowTests/OperatorTests/MathTests.swift

@@ -340,6 +340,12 @@ final class MathOperatorTests: XCTestCase {
         assertEqual(y0, expectedY0, accuracy: 0.0001)
         assertEqual(y1, expectedY1, accuracy: 0.0001)
         assertEqual(y2, expectedY2, accuracy: 0.0001)
+
+        let xSmall = Tensor<Float>([
+            -301.9475, -265.2244, -275.77475, -235.28029, -277.2509, -396.6921, -400.01385])
+        let ySmall = xSmall.logSumExp()
+        let expectedYSmall = Tensor<Float>(-235.28029)
+        assertEqual(ySmall, expectedYSmall, accuracy: 0.0001)
     }
 
     func testMoments() {