[AutoDiff] Add 'zeroTangentVector' property to 'Differentiable' protocol.

[rxwei]

Zero tangent vector is necessary for optimizations on models with an array of parameters, especially for optimizers that iterates over parameters using key paths. The current implementation of some key-path-based optimizers is wrong in that it won't work with models that contain an array of parameters (tangent vectors like infinityNorm are initialized as .zero).

An earlier version of these optimizer using the deprecated AllDifferentiableVariables property would give the correct results, but would be heavyweight and inefficient because they'd need to 1. add a constraint TangentVector == AllDifferentiableVariables to optimizers, and 2. make a copy of all parameters and resetting them to .zero. Since we are deprecating AllDifferentiableVariables, this is not the right direction.

This problem also means that our Differentiable abstraction needs to provide a general mechanism of obtaining a zero tangent vector at a certain instance. Hence we add a zeroTangentVector property to the Differentiable protocol.

Zero tangent vectors do not have a canonical mathematical definition, but makes sense for Differentiable in the standard library because Swift does not have dependent types and thus cannot have a TangentVector that depends on a point on a differentiable manifold. Manopt also has an API, M.zerovec(x), that creates a zero tangent vector at a point (see their API doc here).

Adding zeroTangentVector will make it possible to deprecate AllDifferentiableVariables completely, because currently some fast.ai notebooks depend on initializing parameter gradients using AllDifferentiableVariables.

The new Differentiable protocol looks like the following. The design overview has been updated to reflect this change.

protocol Differentiable {
    /// A type representing the differentiable value’s derivatives.
    /// Mathematically, this is equivalent to the tangent bundle of the
    /// differentiable manifold represented by the differentiable type.
    associatedtype TangentVector: Differentiable & AdditiveArithmetic

    /// Moves `self` along the given direction. In Riemannian geometry,
    /// this is equivalent to exponential map, which moves `self` on the
    /// geodesic surface along the given tangent vector.
    mutating func move(along direction: TangentVector)

    /// A tangent vector such that `move(along: zeroTangentVector)`
    /// will not modify `self`.
    /// - Note: `zeroTangentVector` can be `TangentVector.zero` in most cases,
    ///   but types whose tangent vectors depend on instance properties of
    ///   `self` need to provide a different implementation. For example, an
    ///   array’s zero tangent vector depends on the array’s `count`.
    var zeroTangentVector: TangentVector { get }
}

Resolves TF-708.

[dan-zheng]

Generally LGTM!

(There'll be merge conflicts with my WIP patch to remove AllDifferentiableVariables, though no big deal.)

[dan-zheng]

Big 👍 on these doc comments by the way 🙂

[rxwei]

The API change part of this PR has been implemented as #26828 .

[dan-zheng]

I accidentally deleted tensorflow branch, which closed this PR. That was not intentional, sorry!
It would be nice to protect tensorflow branch against deletion to prevent this from happening again.

[dan-zheng]

Closing this obsolete PR.

Differentiable.zeroTangentVector has been redesigned as Differentiable.zeroTangentVectorInitializer to support efficient lazy zero initialization.

Differentiable.zeroTangentVectorInitializer derived conformances implemented in #31823.