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

Sources/TensorFlow/Operators/Math.swift

@@ -576,6 +576,21 @@ func _vjpLog1p<T: TensorFlowFloatingPoint>(
     (log1p(x), { v in Raw.xdivy(v, 1 + x) })
 }
 
+/// Returns `log(1 - exp(x))` using a numerically stable approach.
+///
+/// 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(_:))
@@ -909,7 +924,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)
@@ -330,30 +337,27 @@ 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),
-        ("testReduction", testReduction),
-        ("testCosineSimilarity", testCosineSimilarity),
-        ("testElu",testElu),
-        ("testGelu", testGelu),
-        ("testArgmax", testArgmax),
         ("testSoftplus", testSoftplus),
         ("testSoftsign", testSoftsign),
         ("testElu",testElu),
+        ("testGelu", testGelu),
         ("testLeakyRelu", testLeakyRelu),
         ("testIsFinite", testIsFinite),
         ("testIsInfinite", testIsInfinite),