Add reflection impl for EuclideanDifferentiable. (#1133)

Author: pschuh

Sources/TensorFlow/CMakeLists.txt

@@ -36,6 +36,7 @@ add_library(TensorFlow SHARED
   Core/TensorShape.swift
   Core/Threading.swift
   Core/Utilities.swift
+  Core/EuclideanDifferentiable.swift
 
   Epochs/Algorithms.swift
   Epochs/Backend.swift
@@ -84,6 +85,7 @@ if(ENABLE_PYTHON_SUPPORT)
 endif()
 target_compile_definitions(TensorFlow PRIVATE
   USING_X10_BACKEND
+  $<$<BOOL:${TENSORFLOW_USE_STANDARD_TOOLCHAIN}>:TENSORFLOW_USE_STANDARD_TOOLCHAIN>
   DEFAULT_BACKEND_EAGER)
 target_link_libraries(TensorFlow PRIVATE
   CX10

Sources/TensorFlow/Core/EuclideanDifferentiable.swift

@@ -0,0 +1,113 @@
+// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import _Differentiation
+
+#if TENSORFLOW_USE_STANDARD_TOOLCHAIN
+@_spi(Reflection) import Swift
+
+func listFields<Root>(of type: Root.Type) -> [(String, PartialKeyPath<Root>)] {
+  var out = [(String, PartialKeyPath<Root>)]()
+  _forEachFieldWithKeyPath(of: type, options: .ignoreUnknown) { name, kp in
+    out.append((String(validatingUTF8: name)!, kp))
+    return true
+  }
+  return out
+}
+
+extension Differentiable {
+  static var differentiableFields: [(String, PartialKeyPath<Self>, PartialKeyPath<TangentVector>)]
+  {
+    let tangentFields = listFields(of: TangentVector.self)
+    var i = 0
+    var out = [(String, PartialKeyPath<Self>, PartialKeyPath<TangentVector>)]()
+    _forEachFieldWithKeyPath(of: Self.self, options: .ignoreUnknown) { cname, kp in
+      if i >= tangentFields.count { return false }
+      let name = String(validatingUTF8: cname)!
+      if tangentFields[i].0 == name {
+        out.append((name, kp, tangentFields[i].1))
+        i += 1
+      }
+      return true
+    }
+    return out
+  }
+}
+
+public protocol _EuclideanDifferentiable {
+  static func _copyWeightsToTangentVector<Root: Differentiable>(
+    _ base: Root, _ out: inout Root.TangentVector,
+    _ keyPathBase: PartialKeyPath<Root>,
+    _ keyPathOut: PartialKeyPath<Root.TangentVector>
+  )
+}
+
+public protocol EuclideanDifferentiable: _EuclideanDifferentiable & Differentiable {
+  var differentiableVectorView: TangentVector { get }
+  func _copyWeightsToTangentVector(_ out: inout TangentVector)
+}
+
+extension EuclideanDifferentiable where TangentVector == Self {
+  public var differentiableVectorView: TangentVector { _read { yield self } }
+  public func _copyWeightsToTangentVector(_ out: inout TangentVector) {
+    out = differentiableVectorView
+  }
+}
+
+extension EuclideanDifferentiable {
+  public static func _copyWeightsToTangentVector<Root: Differentiable>(
+    _ base: Root, _ out: inout Root.TangentVector,
+    _ keyPathBase: PartialKeyPath<Root>,
+    _ keyPathOut: PartialKeyPath<Root.TangentVector>
+  ) {
+    guard let keyPathBase = keyPathBase as? WritableKeyPath<Root, Self>,
+      let keyPathOut = keyPathOut as? WritableKeyPath<Root.TangentVector, Self.TangentVector>
+    else {
+      fatalError("Failure to build differentiableVectorView via reflection: \(Self.self)")
+    }
+    base[keyPath: keyPathBase]._copyWeightsToTangentVector(&out[keyPath: keyPathOut])
+  }
+  public var differentiableVectorView: TangentVector {
+    var out = TangentVector.zero
+    _copyWeightsToTangentVector(&out)
+    return out
+  }
+  public func _copyWeightsToTangentVector(_ out: inout TangentVector) {
+    for (_, keyPathBase, keyPathOut) in Self.differentiableFields {
+      let valueType = type(of: keyPathBase).valueType
+      if let valueType = valueType as? _EuclideanDifferentiable.Type {
+        valueType._copyWeightsToTangentVector(self, &out, keyPathBase, keyPathOut)
+      } else {
+        fatalError("Failure to build differentiableVectorView via reflection: \(valueType)")
+      }
+    }
+  }
+}
+
+extension Float: EuclideanDifferentiable {}
+extension Double: EuclideanDifferentiable {}
+
+extension Array: EuclideanDifferentiable & _EuclideanDifferentiable
+where Element: EuclideanDifferentiable {
+  public func _copyWeightsToTangentVector(_ out: inout TangentVector) {
+    out = Array.DifferentiableView.TangentVector(self.map { $0.differentiableVectorView })
+  }
+}
+extension Array.DifferentiableView: EuclideanDifferentiable & _EuclideanDifferentiable
+where Element: EuclideanDifferentiable {
+  public func _copyWeightsToTangentVector(_ out: inout TangentVector) {
+    out = Array.DifferentiableView.TangentVector(self.base.map { $0.differentiableVectorView })
+  }
+}
+#endif

Sources/TensorFlow/StdlibExtensions.swift

@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import _Differentiation
+@_exported import _Differentiation
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 import Numerics
 #endif
 
+#if !TENSORFLOW_USE_STANDARD_TOOLCHAIN
 // MARK: - Array extensions
 
 extension Array: ElementaryFunctions where Element: ElementaryFunctions {
@@ -107,6 +108,7 @@ extension Array: ElementaryFunctions where Element: ElementaryFunctions {
   /// For complex types, there is a branch cut along the negative real axis.
   public static func root(_ x: Self, _ n: Int) -> Self { x.map { Element.root($0, n) } }
 }
+#endif
 
 // MARK: - Array derivative extensions
 
@@ -116,47 +118,48 @@ 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 { .init(x.map(Element.sqrt)) }
 
   /// 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 { .init(x.map(Element.cos)) }
 
   /// 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 { .init(x.map(Element.sin)) }
 
   /// 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 { .init(x.map(Element.tan)) }
 
   /// 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 { .init(x.map(Element.acos)) }
 
   /// 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 { .init(x.map(Element.asin)) }
 
   /// 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 { .init(x.map(Element.atan)) }
 
   /// The hyperbolic cosine of `x`.
-  public static func cosh(_ x: Self) -> Self { .init(Array.cosh(x.base)) }
+  public static func cosh(_ x: Self) -> Self { .init(x.map(Element.cosh)) }
 
   /// The hyperbolic sine of `x`.
-  public static func sinh(_ x: Self) -> Self { .init(Array.sinh(x.base)) }
+  public static func sinh(_ x: Self) -> Self { .init(x.map(Element.sinh)) }
 
   /// The hyperbolic tangent of `x`.
-  public static func tanh(_ x: Self) -> Self { .init(Array.tanh(x.base)) }
+  public static func tanh(_ x: Self) -> Self { .init(x.map(Element.tanh)) }
 
   /// 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 { .init(x.map(Element.acosh)) }
 
   /// 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 { .init(x.map(Element.asinh)) }
 
   /// 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 { .init(x.map(Element.atanh)) }
 
   /// 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 { .init(x.map(Element.exp)) }
 
+#if !TENSORFLOW_USE_STANDARD_TOOLCHAIN
   /// Two raised to to power `x`.
   public static func exp2(_ x: Self) -> Self { .init(Array.exp2(x.base)) }
 
@@ -165,36 +168,51 @@ where Element: Differentiable & ElementaryFunctions {
 
   /// `exp(x) - 1` evaluated so as to preserve accuracy close to zero.
   public static func expm1(_ x: Self) -> Self { .init(Array.expm1(x.base)) }
+#else
+
+  /// `exp(x) - 1` evaluated so as to preserve accuracy close to zero.
+  public static func expMinusOne(_ x: Self) -> Self { .init(x.map(Element.expMinusOne)) }
+#endif
 
   /// The natural logarithm of `x`.
-  public static func log(_ x: Self) -> Self { .init(Array.log(x.base)) }
+  public static func log(_ x: Self) -> Self { .init(x.map { Element.exp($0) }) }
 
+#if !TENSORFLOW_USE_STANDARD_TOOLCHAIN
   /// The base-two logarithm of `x`.
   public static func log2(_ x: Self) -> Self { .init(Array.log2(x.base)) }
 
   /// The base-ten logarithm of `x`.
   public static func log10(_ x: Self) -> Self { .init(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 {
+    .init(Array.log1p(x.base))
+  }
+#else
+
+  /// The natural logarithm of `x + 1` to preserve accuracy close to zero.
+  public static func log(onePlus x: Self) -> Self {
+    .init(x.map { Element.log(onePlus: $0) })
+  }
+#endif
 
   /// `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 { .init(zip(x, y).map(Element.pow)) }
 
   /// `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 { .init(x.map { Element.pow($0, 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 { .init(x.map { Element.root($0, n) }) }
 }
 
 extension Array.DifferentiableView:
@@ -226,6 +244,7 @@ where Element: Differentiable {
   public init() { self.init(.init()) }
 }
 
+#if !TENSORFLOW_USE_STANDARD_TOOLCHAIN
 extension Array.DifferentiableView: VectorProtocol
 where Element: Differentiable & VectorProtocol {
   public typealias VectorSpaceScalar = Element.VectorSpaceScalar
@@ -282,6 +301,7 @@ where Element: Differentiable & PointwiseMultiplicative {
     }
   }
 }
+#endif
 
 extension Collection {
   /// Returns the `n`th position in `self`.