[Tensor] Add a VJP for 'init(_:)' (scalars) and 'init(shape:scalars:)'. (#512)

Add a VJP for the `init(_:)` which takes an array of scalars and `init(shape:scalars:)` which takes a shape and an array of scalars. Both VJPs are merely returning the original value and a transpose, because `init(_:)` and `init(shape:scalars:)` are linear.

Resolves #510 and #511.

Author: rxwei

Sources/TensorFlow/Core/Tensor.swift

@@ -169,7 +169,7 @@ public extension Tensor {
 }
 
 extension Tensor where Scalar: TensorFlowFloatingPoint {
-    @usableFromInline
+    @inlinable
     func _vjpScalars() -> (value: [Scalar], pullback: (Array<Scalar>.TangentVector) -> Tensor) {
         (value: scalars, pullback: { [shape = self.shape] v in
             Tensor(shape: shape, scalars: v.base)
@@ -200,6 +200,7 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
 public extension Tensor {
     /// Creates a 1D tensor from scalars.
     @inlinable
+    @differentiable(vjp: _vjpInit(_:) where Scalar: TensorFlowFloatingPoint)
     init(_ scalars: [Scalar]) {
         self.init(shape: [scalars.count], scalars: scalars)
     }
@@ -226,6 +227,7 @@ public extension Tensor {
     ///   - scalars: The scalar contents of the tensor.
     /// - Precondition: The product of the dimensions of the shape must equal the number of scalars.
     @inlinable
+    @differentiable(vjp: _vjpInit(shape:scalars:) where Scalar: TensorFlowFloatingPoint)
     init(shape: TensorShape, scalars: [Scalar]) {
         precondition(shape.contiguousSize == scalars.count,
             """
@@ -284,6 +286,22 @@ public extension Tensor {
     }
 }
 
+extension Tensor where Scalar: TensorFlowFloatingPoint {
+    @inlinable
+    static func _vjpInit(_ scalars: [Scalar]) -> (
+        value: Tensor, pullback: (Tensor) -> Array<Scalar>.TangentVector
+    ) {
+        (value: Tensor(scalars), pullback: { v in Array<Scalar>.TangentVector(v.scalars) })
+    }
+
+    @inlinable
+    static func _vjpInit(shape: TensorShape, scalars: [Scalar]) -> (
+        value: Tensor, pullback: (Tensor) -> Array<Scalar>.TangentVector
+    ) {
+        (value: Tensor(scalars), pullback: { v in Array<Scalar>.TangentVector(v.scalars) })
+    }
+}
+
 // Background story on `TensorElementLiteral` and why it's necessary:
 //
 // Very importantly, we want users to be able to implicitly convert an array

Tests/TensorFlowTests/TensorAutoDiffTests.swift

@@ -63,6 +63,20 @@ final class TensorAutoDiffTests: XCTestCase {
         XCTAssertEqual(grad, Tensor([1, 1]))
     }
 
+    func testInitFromScalars() {
+        let grad = gradient(at: [3.0, 4.0]) { x in
+            Tensor(x).sum()
+        }
+        XCTAssertEqual(grad, Array<Double>.TangentVector([1, 1]))
+    }
+
+    func testInitFromScalarsWithShape() {
+        let grad = gradient(at: [3.0, 4.0]) { x in
+            Tensor(shape: [1, 2, 1, 1], scalars: x).sum()
+        }
+        XCTAssertEqual(grad, Array<Double>.TangentVector([1, 1]))
+    }
+
     func testPlus() {
         func f(a: Tensor<Float>, b: Tensor<Float>) -> Tensor<Float> { a + b }
         XCTAssertTrue((Tensor(1), Tensor(1)) == gradient(at: Tensor(0), Tensor(0), in: f))
@@ -516,6 +530,9 @@ final class TensorAutoDiffTests: XCTestCase {
         ("testGenericGrad", testGenericGrad),
         ("testScalarGenericGrad", testScalarGenericGrad),
         ("testScalarized", testScalarized),
+        ("testScalars", testScalars),
+        ("testInitFromScalars", testInitFromScalars),
+        ("testInitFromScalarsWithShape", testInitFromScalarsWithShape),
         ("testPlus", testPlus),
         ("testSubtract", testSubtract),
         ("testMultiply", testMultiply),