Zero tangent vector initializer

zeroTangentVectorInitializer.swift

import TensorFlow

protocol Differentiable {
    ...
    var zeroTangentVectorInitializer: () -> Self { get }
}

extension Tensor where Scalar: TensorFlowFloatingPoint {
    var zeroTangentVectorInitializer: () -> Self {
        { [shape = self.shape] in Tensor(zeros: shape) }
    }
}

struct Foo : Differentiable {
    var x: Tensor<Float>
    var y: Tensor<Float>
    @noDerivative var flag: Bool

    var zeroTangentVectorInitializer: () -> TangentVector {
        let xTanInit = x.zeroTangentVectorInitializer
        let yTanInit = y.zeroTangentVectorInitializer
        return { TangentVector(x: xTanInit(), y: yTanInit()) }
    }
}

func foo(_ a: Foo) -> Tensor<Float> {
    a.y
}

@differentiating(foo)
func foo_vjp(_ a: Foo) -> (value: Tensor<Float>, pullback: (Tensor<Float>) -> Foo.TangentVector) {
    let xTanInit = a.x.zeroTangentVectorInitializer
    return (value: foo(a), pullback: { v in Foo.TangentVector(x: xTanInit(), y: v) })
}