Make 'Module' refine 'EuclideanDifferentiable'. (#491)

`EuclideanDifferentiable`, introduced in swiftlang/swift#26827 and swiftlang/swift#26867, provides a property, `differentiableVectorView`, to project the vector space component of a differentiable struct as a value of  `TangentVector` type. This allows us to express optimization techniques that require the parameter space to be a vector space, e.g. weight decay:
```swift
let 𝛁L: Model.TangentVector = ...
model.move(along: -η * 𝛁L - η * λ * model.differentiableVectorView))
```

This patch makes `Module` refine `EuclideanDifferentiable` and makes all layer combinators (e.g. `Sequential`, `RNN`) be `EuclideanDifferentiable`. All `Module`s and `Layers`s will now have property `differentiableVectorView`.

Resolves #456.

Author: rxwei

Sources/TensorFlow/Core/Tensor.swift

@@ -573,6 +573,6 @@ extension Tensor: PointwiseMultiplicative where Scalar: Numeric {
 // Differentiable
 //===------------------------------------------------------------------------------------------===//
 
-extension Tensor: Differentiable where Scalar: TensorFlowFloatingPoint {
+extension Tensor: Differentiable & EuclideanDifferentiable where Scalar: TensorFlowFloatingPoint {
     public typealias TangentVector = Tensor
 }

Sources/TensorFlow/Layer.swift

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-public protocol Module: Differentiable, KeyPathIterable
+public protocol Module: EuclideanDifferentiable, KeyPathIterable
     where TangentVector: VectorProtocol & ElementaryFunctions &
                          PointwiseMultiplicative & KeyPathIterable {
     /// The input type of the layer.
@@ -52,9 +52,10 @@ public extension Layer {
 }
 
 /// An empty struct representing empty `TangentVector`s for parameterless layers.
-public struct EmptyTangentVector: Differentiable, VectorProtocol, ElementaryFunctions,
+public struct EmptyTangentVector: EuclideanDifferentiable, VectorProtocol, ElementaryFunctions,
                                   PointwiseMultiplicative, KeyPathIterable {
     public typealias VectorSpaceScalar = Float
+    public typealias TangentVector = Self
 
     public func adding(_ x: Float) -> EmptyTangentVector { self }
     public mutating func add(_ x: Float) {}
@@ -67,13 +68,14 @@ public struct EmptyTangentVector: Differentiable, VectorProtocol, ElementaryFunc
 /// A parameterless neural network layer.
 ///
 /// The `TangentVector` of parameterless layers is always `EmptyTangentVector`.
-public protocol ParameterlessLayer: Layer {
+public protocol ParameterlessLayer: Layer where TangentVector == EmptyTangentVector {
     @differentiable
     func callAsFunction(_ input: Input) -> Output
 }
 
 public extension ParameterlessLayer {
     mutating func move(along direction: EmptyTangentVector) {}
+    var differentiableVectorView: EmptyTangentVector { EmptyTangentVector() }
 }
 
 public extension Layer {

Sources/TensorFlow/Layers/Recurrent.swift

@@ -26,6 +26,9 @@ public struct RNNCellInput<Input: Differentiable, State: Differentiable>: Differ
     }
 }
 
+extension RNNCellInput: EuclideanDifferentiable
+    where Input: EuclideanDifferentiable, State: EuclideanDifferentiable {}
+
 /// An output to a recurrent neural network.
 public struct RNNCellOutput<Output: Differentiable, State: Differentiable>: Differentiable {
     /// The output at the current time step.
@@ -40,6 +43,9 @@ public struct RNNCellOutput<Output: Differentiable, State: Differentiable>: Diff
     }
 }
 
+extension RNNCellOutput: EuclideanDifferentiable
+    where Output: EuclideanDifferentiable, State: EuclideanDifferentiable {}
+
 /// A recurrent neural network cell.
 public protocol RNNCell: Layer
     where Input == RNNCellInput<TimeStepInput, State>, 

Tests/TensorFlowTests/LayerTests.swift

@@ -243,7 +243,6 @@ final class LayerTests: XCTestCase {
         XCTAssertEqual(output, expected)
     }
 
-
     func testZeroPadding1D() {
         let input = Tensor<Float>([0.0, 1.0, 2.0])
         let layer = ZeroPadding1D<Float>(padding: 2)