Added support for the 'sign' op and its VJP. (#144)

Author: eaplatanios

Sources/TensorFlow/Operators/Math.swift

@@ -597,6 +597,19 @@ internal func _vjpFloor<T: TensorFlowFloatingPoint>(
     return (floor(x), { _ in Tensor(0).broadcasted(like: x) })
 }
 
+@inlinable
+@differentiable(vjp: _vjpSign(_:) where T: TensorFlowFloatingPoint)
+public func sign<T: Numeric>(_ x: Tensor<T>) -> Tensor<T> {
+    return Raw.sign(x)
+}
+
+@inlinable
+internal func _vjpSign<T: TensorFlowFloatingPoint>(
+    _ x: Tensor<T>
+) -> (Tensor<T>, (Tensor<T>) -> Tensor<T>) {
+    return (sign(x), { v in Tensor<T>(zerosLike: x) })
+}
+
 /// Computes the sigmoid of the specified tensor element-wise.
 /// Specifically, computes `1 / (1 + exp(-x))`.
 @inlinable

Tests/TensorFlowTests/OperatorTests/MathTests.swift

@@ -16,6 +16,12 @@ import XCTest
 @testable import TensorFlow
 
 final class MathOperatorTests: XCTestCase {
+    func testSign() {
+        let x = Tensor<Float>([[1, 2, -3, 4, 5], [1, 2, 3, 4, -5]])
+        let y = sign(x)
+        XCTAssertEqual(y, Tensor<Float>([[1, 1, -1, 1, 1], [1, 1, 1, 1, -1]]))
+    }
+
     func testReduction() {
         // 2 x 5
         let x = Tensor<Float>([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]])
@@ -211,6 +217,7 @@ final class MathOperatorTests: XCTestCase {
 	}
 
     static var allTests = [
+        ("testSign", testSign),
         ("testReduction", testReduction),
         ("testArgmax", testArgmax),
         ("testCeilAndFloor", testCeilAndFloor),