Clean up.

There are two remaining test failures:

```
$ swift test --filter OptimizerTests
Tests/TensorFlowTests/OptimizerTests.swift:117: error: -[TensorFlowTests.OptimizerTests testAdaMax] : XCTAssertTrue failed
Tests/TensorFlowTests/OptimizerTests.swift:123: error: -[TensorFlowTests.OptimizerTests testAMSGrad] : XCTAssertTrue failed
```

Author: dan-zheng

Sources/TensorFlow/Layers/Dense.swift

@@ -40,19 +40,20 @@ public struct Dense<Scalar: TensorFlowFloatingPoint>: Layer {
 
   /// The bias vector.
   ///
-  /// - Note: returns `Tensor.zero` if the underlying `optionalBias`  does not exist.
-  //@differentiable
+  /// - Note: Returns `Tensor.zero` if the underlying `optionalBias` does not exist.
+  @differentiable
   public var bias: Tensor<Scalar> {
-    get { optionalBias ?? .zero }
+    get {
+      if let bias = optionalBias {
+        return bias
+      }
+      return .zero
+    }
     set { optionalBias = newValue }
   }
 
   /// Creates an instance from the given weight, optional bias, and activation function.
-  ///
-  /// - Note: currently, `weight` is the only differentiability parameter. `bias` can be made a
-  ///   differentiability parameter after `Optional` conditionally conforms to `Differentiable`:
-  ///   TF-499.
-  @differentiable(wrt: weight)
+  @differentiable(wrt: (weight, bias))
   public init(
     weight: Tensor<Scalar>,
     bias: Tensor<Scalar>? = nil,
@@ -67,18 +68,6 @@ public struct Dense<Scalar: TensorFlowFloatingPoint>: Layer {
     self.batched = weight.rank == 3
   }
 
-  // TODO(TF-433): Remove custom derivative after `try_apply` differentiation is supported.
-  @derivative(of: init, wrt: weight)
-  @usableFromInline
-  static func vjpInit(
-    weight: Tensor<Scalar>,
-    bias: Tensor<Scalar>? = nil,
-    activation: @escaping Activation
-  ) -> (value: Self, pullback: (TangentVector) -> Tensor<Scalar>) {
-    let value = Dense(weight: weight, bias: bias, activation: activation)
-    return (value, { v in v.weight })
-  }
-
   /// Returns the output obtained from applying the layer to the given input.
   ///
   /// - Parameter input: The input to the layer.
@@ -118,7 +107,6 @@ extension Dense {
     weightInitializer: ParameterInitializer<Scalar> = glorotUniform(),
     biasInitializer: ParameterInitializer<Scalar> = zeros()
   ) {
-    print("Init OLD")
     self.init(
       weight: weightInitializer([inputSize, outputSize]),
       bias: useBias ? biasInitializer([outputSize]) : nil,
@@ -142,7 +130,6 @@ extension Dense {
     weightInitializer: ParameterInitializer<Scalar> = glorotUniform(),
     biasInitializer: ParameterInitializer<Scalar>? = nil
   ) {
-    print("Init NEW")
     self.init(
       weight: weightInitializer([inputSize, outputSize]),
       bias: biasInitializer?([outputSize]),
@@ -168,18 +155,13 @@ extension Dense.TangentVector {
   }
 }
 
-/* extension Optional : KeyPathIterable {
+extension Optional: KeyPathIterable {
   public var allKeyPaths: [PartialKeyPath<Self>] {
     if self != nil {
-      return [ \Optional.unsafelyUnwrapped ]
+      return [\.!]
     }
     return []
   }
 
   public typealias AllKeyPaths = [PartialKeyPath<Self>]
 }
-
-extension Optional.TangentVector : KeyPathIterable
-{
-  
-}*/

Sources/TensorFlow/StdlibExtensions.swift

@@ -111,85 +111,85 @@ where Element: Differentiable & ElementaryFunctions {
   ///
   /// For real types, if `x` is negative the result is `.nan`. For complex
   /// types there is a branch cut on the negative real axis.
-  public static func sqrt(_ x: Self) -> Self { .init(Array.sqrt(x.base)) }
+  public static func sqrt(_ x: Self) -> Self { Self(Array.sqrt(x.base)) }
 
   /// The cosine of `x`, interpreted as an angle in radians.
-  public static func cos(_ x: Self) -> Self { .init(Array.cos(x.base)) }
+  public static func cos(_ x: Self) -> Self { Self(Array.cos(x.base)) }
 
   /// The sine of `x`, interpreted as an angle in radians.
-  public static func sin(_ x: Self) -> Self { .init(Array.sin(x.base)) }
+  public static func sin(_ x: Self) -> Self { Self(Array.sin(x.base)) }
 
   /// The tangent of `x`, interpreted as an angle in radians.
-  public static func tan(_ x: Self) -> Self { .init(Array.tan(x.base)) }
+  public static func tan(_ x: Self) -> Self { Self(Array.tan(x.base)) }
 
   /// The inverse cosine of `x` in radians.
-  public static func acos(_ x: Self) -> Self { .init(Array.acos(x.base)) }
+  public static func acos(_ x: Self) -> Self { Self(Array.acos(x.base)) }
 
   /// The inverse sine of `x` in radians.
-  public static func asin(_ x: Self) -> Self { .init(Array.asin(x.base)) }
+  public static func asin(_ x: Self) -> Self { Self(Array.asin(x.base)) }
 
   /// The inverse tangent of `x` in radians.
-  public static func atan(_ x: Self) -> Self { .init(Array.atan(x.base)) }
+  public static func atan(_ x: Self) -> Self { Self(Array.atan(x.base)) }
 
   /// The hyperbolic cosine of `x`.
-  public static func cosh(_ x: Self) -> Self { .init(Array.cosh(x.base)) }
+  public static func cosh(_ x: Self) -> Self { Self(Array.cosh(x.base)) }
 
   /// The hyperbolic sine of `x`.
-  public static func sinh(_ x: Self) -> Self { .init(Array.sinh(x.base)) }
+  public static func sinh(_ x: Self) -> Self { Self(Array.sinh(x.base)) }
 
   /// The hyperbolic tangent of `x`.
-  public static func tanh(_ x: Self) -> Self { .init(Array.tanh(x.base)) }
+  public static func tanh(_ x: Self) -> Self { Self(Array.tanh(x.base)) }
 
   /// The inverse hyperbolic cosine of `x`.
-  public static func acosh(_ x: Self) -> Self { .init(Array.acosh(x.base)) }
+  public static func acosh(_ x: Self) -> Self { Self(Array.acosh(x.base)) }
 
   /// The inverse hyperbolic sine of `x`.
-  public static func asinh(_ x: Self) -> Self { .init(Array.asinh(x.base)) }
+  public static func asinh(_ x: Self) -> Self { Self(Array.asinh(x.base)) }
 
   /// The inverse hyperbolic tangent of `x`.
-  public static func atanh(_ x: Self) -> Self { .init(Array.atanh(x.base)) }
+  public static func atanh(_ x: Self) -> Self { Self(Array.atanh(x.base)) }
 
   /// The exponential function applied to `x`, or `e**x`.
-  public static func exp(_ x: Self) -> Self { .init(Array.exp(x.base)) }
+  public static func exp(_ x: Self) -> Self { Self(Array.exp(x.base)) }
 
   /// Two raised to to power `x`.
-  public static func exp2(_ x: Self) -> Self { .init(Array.exp2(x.base)) }
+  public static func exp2(_ x: Self) -> Self { Self(Array.exp2(x.base)) }
 
   /// Ten raised to to power `x`.
-  public static func exp10(_ x: Self) -> Self { .init(Array.exp10(x.base)) }
+  public static func exp10(_ x: Self) -> Self { Self(Array.exp10(x.base)) }
 
   /// `exp(x) - 1` evaluated so as to preserve accuracy close to zero.
-  public static func expm1(_ x: Self) -> Self { .init(Array.expm1(x.base)) }
+  public static func expm1(_ x: Self) -> Self { Self(Array.expm1(x.base)) }
 
   /// The natural logarithm of `x`.
-  public static func log(_ x: Self) -> Self { .init(Array.log(x.base)) }
+  public static func log(_ x: Self) -> Self { Self(Array.log(x.base)) }
 
   /// The base-two logarithm of `x`.
-  public static func log2(_ x: Self) -> Self { .init(Array.log2(x.base)) }
+  public static func log2(_ x: Self) -> Self { Self(Array.log2(x.base)) }
 
   /// The base-ten logarithm of `x`.
-  public static func log10(_ x: Self) -> Self { .init(Array.log10(x.base)) }
+  public static func log10(_ x: Self) -> Self { Self(Array.log10(x.base)) }
 
   /// `log(1 + x)` evaluated so as to preserve accuracy close to zero.
-  public static func log1p(_ x: Self) -> Self { .init(Array.log1p(x.base)) }
+  public static func log1p(_ x: Self) -> Self { Self(Array.log1p(x.base)) }
 
   /// `exp(y log(x))` computed without loss of intermediate precision.
   ///
   /// For real types, if `x` is negative the result is NaN, even if `y` has
   /// an integral value. For complex types, there is a branch cut on the
   /// negative real axis.
-  public static func pow(_ x: Self, _ y: Self) -> Self { .init(Array.pow(x.base, y.base)) }
+  public static func pow(_ x: Self, _ y: Self) -> Self { Self(Array.pow(x.base, y.base)) }
 
   /// `x` raised to the `n`th power.
   ///
   /// The product of `n` copies of `x`.
-  public static func pow(_ x: Self, _ n: Int) -> Self { .init(Array.pow(x.base, n)) }
+  public static func pow(_ x: Self, _ n: Int) -> Self { Self(Array.pow(x.base, n)) }
 
   /// The `n`th root of `x`.
   ///
   /// For real types, if `x` is negative and `n` is even, the result is NaN.
   /// For complex types, there is a branch cut along the negative real axis.
-  public static func root(_ x: Self, _ n: Int) -> Self { .init(Array.root(x.base, n)) }
+  public static func root(_ x: Self, _ n: Int) -> Self { Self(Array.root(x.base, n)) }
 }
 
 extension Array.DifferentiableView:
@@ -226,7 +226,7 @@ where Element: Differentiable & VectorProtocol {
   public typealias VectorSpaceScalar = Element.VectorSpaceScalar
 
   public func adding(_ x: Element.VectorSpaceScalar) -> Array<Element>.DifferentiableView {
-    .init(map { $0.adding(x) })
+    Self(map { $0.adding(x) })
   }
 
   public mutating func add(_ x: Element.VectorSpaceScalar) {
@@ -236,7 +236,7 @@ where Element: Differentiable & VectorProtocol {
   }
 
   public func subtracting(_ x: Element.VectorSpaceScalar) -> Array<Element>.DifferentiableView {
-    .init(map { $0.subtracting(x) })
+    Self(map { $0.subtracting(x) })
   }
 
   public mutating func subtract(_ x: Element.VectorSpaceScalar) {
@@ -246,7 +246,7 @@ where Element: Differentiable & VectorProtocol {
   }
 
   public func scaled(by scale: Element.VectorSpaceScalar) -> Self {
-    .init(map { $0.scaled(by: scale) })
+    Self(map { $0.scaled(by: scale) })
   }
 
   public mutating func scale(by scale: Element.VectorSpaceScalar) {
@@ -263,11 +263,11 @@ where Element: Differentiable & PointwiseMultiplicative {
     fatalError("One is not array-representable")
   }
 
-  public var reciprocal: Self { .init(map { $0.reciprocal }) }
+  public var reciprocal: Self { Self(map { $0.reciprocal }) }
 
   public static func .* (lhs: Self, rhs: Self) -> Self {
     precondition(lhs.count == rhs.count, "Count mismatch: \(lhs.count) and \(rhs.count)")
-    return .init(zip(lhs, rhs).map(.*))
+    return Self(zip(lhs, rhs).map(.*))
   }
 
   public static func .*= (lhs: inout Self, rhs: Self) {
@@ -294,94 +294,103 @@ where Wrapped.TangentVector: ElementaryFunctions {
   ///
   /// For real types, if `x` is negative the result is `.nan`. For complex
   /// types there is a branch cut on the negative real axis.
-  public static func sqrt(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func sqrt(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.sqrt)) }
 
   /// The cosine of `x`, interpreted as an angle in radians.
-  public static func cos(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func cos(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.cos)) }
 
   /// The sine of `x`, interpreted as an angle in radians.
-  public static func sin(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func sin(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.sin)) }
 
   /// The tangent of `x`, interpreted as an angle in radians.
-  public static func tan(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func tan(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.tan)) }
 
   /// The inverse cosine of `x` in radians.
-  public static func acos(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func acos(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.acos)) }
 
   /// The inverse sine of `x` in radians.
-  public static func asin(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func asin(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.asin)) }
 
   /// The inverse tangent of `x` in radians.
-  public static func atan(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func atan(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.atan)) }
 
   /// The hyperbolic cosine of `x`.
-  public static func cosh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func cosh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.cosh)) }
 
   /// The hyperbolic sine of `x`.
-  public static func sinh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func sinh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.sinh)) }
 
   /// The hyperbolic tangent of `x`.
-  public static func tanh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func tanh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.tanh)) }
 
   /// The inverse hyperbolic cosine of `x`.
-  public static func acosh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func acosh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.acosh)) }
 
   /// The inverse hyperbolic sine of `x`.
-  public static func asinh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func asinh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.asinh)) }
 
   /// The inverse hyperbolic tangent of `x`.
-  public static func atanh(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func atanh(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.atanh)) }
 
   /// The exponential function applied to `x`, or `e**x`.
-  public static func exp(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func exp(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.exp)) }
 
   /// Two raised to to power `x`.
-  public static func exp2(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func exp2(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.exp2)) }
 
   /// Ten raised to to power `x`.
-  public static func exp10(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func exp10(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.exp10)) }
 
   /// `exp(x) - 1` evaluated so as to preserve accuracy close to zero.
-  public static func expm1(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func expm1(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.expm1)) }
 
   /// The natural logarithm of `x`.
-  public static func log(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func log(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.log)) }
 
   /// The base-two logarithm of `x`.
-  public static func log2(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func log2(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.log2)) }
 
   /// The base-ten logarithm of `x`.
-  public static func log10(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func log10(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.log10)) }
 
   /// `log(1 + x)` evaluated so as to preserve accuracy close to zero.
-  public static func log1p(_ x: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func log1p(_ x: Self) -> Self { Self(x.value.map(Wrapped.TangentVector.log1p)) }
 
   /// `exp(y log(x))` computed without loss of intermediate precision.
   ///
   /// For real types, if `x` is negative the result is NaN, even if `y` has
   /// an integral value. For complex types, there is a branch cut on the
   /// negative real axis.
-  public static func pow(_ x: Self, _ y: Self) -> Self { .init(x.value.map(Wrapped.TangentVector.sqrt)) }
+  public static func pow(_ x: Self, _ y: Self) -> Self {
+    switch (x.value, y.value) {
+    case let (x?, y?): return Self(Wrapped.TangentVector.pow(x, y))
+    default: return Self(nil)
+    }
+  }
 
   /// `x` raised to the `n`th power.
   ///
   /// The product of `n` copies of `x`.
-    public static func pow(_ x: Self, _ n: Int) -> Self { .init(x.value.map({ x in Wrapped.TangentVector.pow(x, n)})) }
+  public static func pow(_ x: Self, _ n: Int) -> Self {
+    Self(x.value.map({ x in Wrapped.TangentVector.pow(x, n) }))
+  }
 
   /// The `n`th root of `x`.
   ///
   /// For real types, if `x` is negative and `n` is even, the result is NaN.
   /// For complex types, there is a branch cut along the negative real axis.
-  public static func root(_ x: Self, _ n: Int) -> Self { .init(x.value.map({ x in Wrapped.TangentVector.root(x, n)})) }
+  public static func root(_ x: Self, _ n: Int) -> Self {
+    Self(x.value.map({ x in Wrapped.TangentVector.root(x, n) }))
+  }
 }
 
 extension Optional.TangentVector: PointwiseMultiplicative
 where Wrapped.TangentVector: PointwiseMultiplicative {
   public static var one: Self {
-    .init(Wrapped.TangentVector.one)
+    Self(Wrapped.TangentVector.one)
   }
 
-  public var reciprocal: Self { .init(value.map { $0.reciprocal }) }
+  public var reciprocal: Self { Self(value.map { $0.reciprocal }) }
 
   public static func .* (lhs: Self, rhs: Self) -> Self {
     switch (lhs.value, rhs.value) {
@@ -399,15 +408,17 @@ extension Optional.TangentVector: VectorProtocol
 where Wrapped.TangentVector: VectorProtocol {
   public typealias VectorSpaceScalar = Wrapped.TangentVector.VectorSpaceScalar
 
-  public func adding(_ x: VectorSpaceScalar) -> Self { .init(value.map { $0.adding(x) }) }
+  public func adding(_ x: VectorSpaceScalar) -> Self { Self(value.map { $0.adding(x) }) }
 
   public mutating func add(_ x: VectorSpaceScalar) { value?.add(x) }
 
-  public func subtracting(_ x: VectorSpaceScalar) -> Self { .init(value.map { $0.subtracting(x) }) }
+  public func subtracting(_ x: VectorSpaceScalar) -> Self { Self(value.map { $0.subtracting(x) }) }
 
   public mutating func subtract(_ x: VectorSpaceScalar) { value?.subtract(x) }
 
-  public func scaled(by scale: VectorSpaceScalar) -> Self { .init(value.map { $0.scaled(by: scale) }) }
+  public func scaled(by scale: VectorSpaceScalar) -> Self {
+    Self(value.map { $0.scaled(by: scale) })
+  }
 
   public mutating func scale(by scale: VectorSpaceScalar) {
     value?.scale(by: scale)

Tests/TensorFlowTests/TrivialModelTests.swift

@@ -51,7 +51,6 @@ final class TrivialModelTests: XCTestCase {
           return meanSquaredError(predicted: ŷ, expected: y)
         }
         optimizer.update(&classifier, along: 𝛁model)
-        dump(𝛁model)
       }
     }
     let ŷ = classifier.inferring(from: x)