[AutoDiff] Define Tensor.zeroTangentVectorInitializer. (#1055)

`Tensor.zeroTangentVectorInitializer` now returns a zero tensor with the same
shape as `self`, instead of a scalar zero tensor.

Author: dan-zheng

Sources/TensorFlow/Core/Tensor.swift

@@ -747,6 +747,11 @@ extension Tensor: PointwiseMultiplicative where Scalar: Numeric {
 
 extension Tensor: Differentiable & EuclideanDifferentiable where Scalar: TensorFlowFloatingPoint {
   public typealias TangentVector = Tensor
+
+  public var zeroTangentVectorInitializer: () -> TangentVector {
+    let shape = self.shape
+    return { Tensor(zeros: shape) }
+  }
 }
 
 //===------------------------------------------------------------------------------------------===//

Tests/TensorFlowTests/TensorTests.swift

@@ -120,6 +120,18 @@ final class TensorTests: XCTestCase {
     )
   }
 
+  func testZeroTangentVectorInitializer() {
+    let shape: TensorShape = [4, 5, 6]
+    let tensor = Tensor<Float>(randomUniform: shape)
+    XCTAssertEqual(tensor.zeroTangentVector, Tensor(zeros: shape))
+
+    struct TensorWrapper: Differentiable {
+      var tensor: Tensor<Float>
+    }
+    let model = TensorWrapper(tensor: tensor)
+    XCTAssertEqual(model.zeroTangentVector, .init(tensor: Tensor(zeros: shape)))
+  }
+
   static var allTests = [
     ("testSimpleCond", testSimpleCond),
     ("testRankGetter", testRankGetter),
@@ -129,5 +141,6 @@ final class TensorTests: XCTestCase {
     ("testTensorShapeCollectionOperations", testTensorShapeCollectionOperations),
     ("testInitShapeScalars", testInitShapeScalars),
     ("testInitShapeScalarsDerivative", testInitShapeScalarsDerivative),
+    ("testZeroTangentVectorInitializer", testZeroTangentVectorInitializer),
   ]
 }