Add derivative for 'padded(forSizes:with:)'. (tensorflow#184)

Author: lakshya-sky

Sources/TensorFlow/Operators/Basic.swift

@@ -674,6 +674,7 @@ extension Tensor where Scalar: TensorFlowFloatingPoint {
 public extension Tensor where Scalar: Numeric {
     /// Returns a padded tensor according to the specified padding sizes.
     @inlinable
+    @differentiable(wrt: self, vjp: _vjpPadded(forSizes:with:) where Scalar: TensorFlowFloatingPoint)
     func padded(forSizes sizes: [(before: Int, after: Int)], with value: Scalar = 0) -> Tensor {
         let paddings = Tensor<Int32>(
             shape: [sizes.count, 2],
@@ -682,6 +683,26 @@ public extension Tensor where Scalar: Numeric {
     }
 }
 
+internal extension Tensor where Scalar: TensorFlowFloatingPoint {
+    @inlinable
+    func _vjpPadded(
+        forSizes sizes: [(before: Int, after: Int)],
+        with value: Scalar
+    ) -> (Tensor, (Tensor) -> Tensor) {
+        let result = padded(forSizes: sizes, with: value)
+        return (result, { [rank = rankTensor, shape = shapeTensor] v in
+            let paddings = Tensor<Int32>(
+                shape: [sizes.count, 2],
+                scalars: sizes.flatMap { [Int32($0.before), Int32($0.after)] })
+            let padBefore = Raw.slice(paddings,
+                begin: Tensor<Int32>([0, 0]),
+                size: Tensor<Int32>(stacking: [rank, Tensor<Int32>(1)]))
+            let begin = padBefore.reshaped(to: [-1])
+            return Raw.slice(v, begin: begin, size: shape)
+        })
+    }
+}
+
 //===------------------------------------------------------------------------------------------===//
 // Indexing and Slicing
 //===------------------------------------------------------------------------------------------===//

Tests/TensorFlowTests/OperatorTests/BasicTests.swift

@@ -32,6 +32,23 @@ final class BasicOperatorTests: XCTestCase {
         XCTAssertEqual(y, Tensor<Float>([3.0, 6.0]))
     }
 
+    func testPadded() {
+        let x = Tensor<Float>(ones: [2, 2])
+        let target = Tensor<Float>([[3, 3, 3], [1, 1, 3], [1, 1, 3]])
+        let paddedTensor = x.padded(forSizes: [(1, 0), (0, 1)], with: 3.0)
+        XCTAssertEqual(paddedTensor, target)
+    }
+
+    func testVJPPadded() {
+        let x = Tensor<Float>(ones: [3, 2])
+        let target = Tensor<Float>([[2, 2], [2, 2], [2, 2]])
+        let grads = x.gradient { a -> Tensor<Float> in
+            let paddedTensor = a.padded(forSizes: [(1, 0), (0, 1)], with: 3.0)
+            return (paddedTensor * paddedTensor).sum()
+        }
+        XCTAssertEqual(grads, target)
+    }
+
     func testElementIndexing() {
         // NOTE: cannot test multiple `Tensor.shape` or `Tensor.scalars` directly
         // until send and receive are implemented (without writing a bunch of mini
@@ -468,6 +485,8 @@ final class BasicOperatorTests: XCTestCase {
 
     static var allTests = [
         ("testGathering", testGathering),
+        ("testPadded", testPadded),
+        ("testVJPPadded", testVJPPadded),
         ("testElementIndexing", testElementIndexing),
         ("testElementIndexingAssignment", testElementIndexingAssignment),
         ("testNestedElementIndexing", testNestedElementIndexing),