Add support for ReLU6 activation (#435)

Author: jon-tow

Sources/TensorFlow/Operators/Math.swift

@@ -1183,6 +1183,20 @@ func _vjpRelu<T: TensorFlowFloatingPoint>(
     (relu(x), { v in Raw.reluGrad(gradients: v, features: x) })
 }
 
+/// Returns a tensor by applying the ReLU6 activation function, namely `min(max(0, x), 6)`.
+@inlinable
+@differentiable(vjp: _vjpRelu6(_:))
+public func relu6<T: TensorFlowFloatingPoint>(_ x: Tensor<T>) -> Tensor<T> {
+    Raw.relu6(features: x)
+}
+
+@inlinable
+func _vjpRelu6<T: TensorFlowFloatingPoint>(
+    _ x: Tensor<T>
+) -> (Tensor<T>, (Tensor<T>) -> Tensor<T>) {
+    (relu6(x), { v in Raw.relu6Grad(gradients: v, features: x)})
+}
+
 /// Returns a tensor by applying the leaky ReLU activation function
 /// to the specified tensor element-wise.
 /// Specifically, computes `max(x, x * alpha)`.

Tests/TensorFlowTests/OperatorTests/MathTests.swift

@@ -153,6 +153,13 @@ final class MathOperatorTests: XCTestCase {
         XCTAssertEqual(y, expectedY)
     }
 
+    func testRelu6() {
+        let x = Tensor<Float>([1.0, -2.0, 3.0, 4.0, 10.0])
+        let y = relu6(x)
+        let expectedY = Tensor<Float>([1.0, 0, 3.0, 4.0, 6.0])
+        XCTAssertEqual(y, expectedY)
+    }
+
     func testLeakyRelu() {
         let x = Tensor<Float>([[-1.0, 2.0, 3.0]])
         let y = leakyRelu(x, alpha: 0.4)
@@ -536,6 +543,7 @@ final class MathOperatorTests: XCTestCase {
         ("testElu",testElu),
         ("testGelu", testGelu),
         ("testRelu", testRelu),
+        ("testRelu6", testRelu6),
         ("testLeakyRelu", testLeakyRelu),
         ("testSelu", testSelu),
         ("testIsFinite", testIsFinite),