Added support for the 'log1mexp' op and its VJP. (#147)

* Enhanced the 'matmul' wrapper so that it matches the behavior of the Python one.

* Added support for the 'log1mexp' op and its VJP.

* Added a test.

* Update Sources/TensorFlow/Operators/Math.swift

Co-Authored-By: Richard Wei <[email protected]>

* Removed the need for a general 'Tensor.withoutDerivative()' as Richard suggested.

* Addressed Richard's feedback.

* Addressed Richard's feedback.

* Added one more tests helper.

* Minor bug fix.

* Added a test for 'log1mexp'.

* Added support for 'softplus' and 'logSigmoid'.

* Minor tweak.

* Fixed some of the tests.

* Made the tests pass.

* Attempt at making 'log1mexp' differentiable.

* Merged upstream changes.

* Enabled the 'logSigmoid' test.

* Style edits.

* Style edits.

* Update Sources/TensorFlow/Operators/Math.swift

Co-Authored-By: Richard Wei <[email protected]>

* Update Sources/TensorFlow/Operators/Math.swift

Co-Authored-By: Richard Wei <[email protected]>

Author: eaplatanios

Sources/TensorFlow/Operators/Math.swift

@@ -579,6 +579,21 @@ func _vjpLog1p<T: TensorFlowFloatingPoint>(
     (log1p(x), { v in Raw.xdivy(v, 1 + x) })
 }
 
+/// Returns `log(1 - exp(x))` using a numerically stable approach.
+///
+/// - Note: The approach is shown in Equation 7 of:
+///   https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
+@inlinable
+@differentiable
+public func log1mexp<T: TensorFlowFloatingPoint>(_ x: Tensor<T>) -> Tensor<T> {
+    let isTooSmall = withoutDerivative(at: x) { x in -x .< T(log(2.0)) }
+    // This `replacing` will ultimately be a no-op because we will not select this code-path 
+    // whenever we use the surrogate `-Tensor(onesLike: x)`.
+    let ones = withoutDerivative(at: x) { x in Tensor(onesLike: x) }
+    let xSafe = x.replacing(with: -ones, where: isTooSmall)
+    return log1p(-exp(xSafe)).replacing(with: log(-expm1(x)), where: isTooSmall)
+}
+
 /// Returns the sine of the specified tensor element-wise.
 @inlinable
 @differentiable(vjp: _vjpSin(_:))
@@ -912,7 +927,7 @@ internal func _vjpSigmoid<T: TensorFlowFloatingPoint>(
 }
 
 /// Returns the log-sigmoid of the specified tensor element-wise. Specifically,
-/// `y = log(1 / (1 + exp(-x)))`. For numerical stability, we use `y = -softplus(-x)`.
+/// `log(1 / (1 + exp(-x)))`. For numerical stability, we use `-softplus(-x)`.
 @inlinable
 @differentiable
 public func logSigmoid<T: TensorFlowFloatingPoint>(_ x: Tensor<T>) -> Tensor<T> {

Tests/TensorFlowTests/OperatorTests/MathTests.swift

@@ -66,6 +66,13 @@ final class MathOperatorTests: XCTestCase {
         assertEqual(y, expectedY, accuracy: 0.0001)
     }
 
+    func testLog1mexp() {
+        let x = Tensor<Float>([-1, -2, -3, -4, -5])
+        let y = log1mexp(x)
+        let expectedY = Tensor<Float>([-0.45868, -0.14541, -0.05107, -0.01849, -0.00676])
+        assertEqual(y, expectedY, accuracy: 0.0001)
+    }
+
     func testExpm1() {
         let x = Tensor<Float>([1, 2, 3, 4, 5])
         let y = expm1(x)
@@ -350,19 +357,20 @@ final class MathOperatorTests: XCTestCase {
     }
 
     func testBroadcastedAddGradient() {
-	  func foo(_ x: Tensor<Float>, _ y: Tensor<Float>) -> Tensor<Float> {
-	    return (x + y).sum()
-	  }
-	  let x = Tensor<Float>(ones: [1, 2, 1, 4])
-	  let y = Tensor<Float>(ones: [4, 1, 3, 1])
-	  let (dx, dy) = gradient(at: x, y, in: foo)
-	  XCTAssertEqual(x.shape, dx.shape)
-	  XCTAssertEqual(y.shape, dy.shape)
-	}
+        func foo(_ x: Tensor<Float>, _ y: Tensor<Float>) -> Tensor<Float> {
+            return (x + y).sum()
+        }
+        let x = Tensor<Float>(ones: [1, 2, 1, 4])
+        let y = Tensor<Float>(ones: [4, 1, 3, 1])
+        let (dx, dy) = gradient(at: x, y, in: foo)
+        XCTAssertEqual(x.shape, dx.shape)
+        XCTAssertEqual(y.shape, dy.shape)
+    }
 
     static var allTests = [
         ("testElementaryFunctions", testElementaryFunctions),
         ("testLog1p", testLog1p),
+        ("testLog1mexp", testLog1mexp),
         ("testExpm1", testExpm1),
         ("testSign", testSign),
         ("testLogSigmoid", testLogSigmoid),