Initial support for release toolchains

philipturner · Jul 3, 2022 · 0cfa5bb · 0cfa5bb
1 parent 81e84dc
commit 0cfa5bb
Show file tree

Hide file tree

Showing 44 changed files with 307 additions and 20 deletions.
diff --git a/Package.swift b/Package.swift
@@ -16,6 +16,25 @@
 // limitations under the License.
 
 import PackageDescription
+import class Foundation.ProcessInfo
+
+var conditionalPackageDependencies: [Package.Dependency] = []
+var conditionalSwiftSettings: [SwiftSetting] = []
+var conditionalTargetDependencies: [Target.Dependency] = []
+
+if ProcessInfo.processInfo.environment["TENSORFLOW_USE_RELEASE_TOOLCHAIN"] != nil {
+  conditionalPackageDependencies += [
+    .package(url: "https://github.com/philipturner/differentiation", .branch("main")),
+    .package(url: "https://github.com/philipturner/swift-reflection-mirror", .branch("main")),
+  ]
+  conditionalSwiftSettings += [
+    .define("TENSORFLOW_USE_RELEASE_TOOLCHAIN"),
+  ]
+  conditionalTargetDependencies += [
+    .product(name: "_Differentiation", package: "differentiation"),
+    .product(name: "ReflectionMirror", package: "swift-reflection-mirror"),
+  ]
+}
 
 let package = Package(
   name: "TensorFlow",
@@ -35,15 +54,15 @@ let package = Package(
       name: "x10_optimizers_tensor_visitor_plan",
       type: .dynamic,
       targets: ["x10_optimizers_tensor_visitor_plan"]),
-    // .library(
-    //   name: "x10_training_loop",
-    //   type: .dynamic,
-    //   targets: ["x10_training_loop"]),
+    .library(
+      name: "x10_training_loop",
+      type: .dynamic,
+      targets: ["x10_training_loop"]),
   ],
   dependencies: [
     .package(url: "https://github.com/apple/swift-numerics", .branch("main")),
     .package(url: "https://github.com/pvieito/PythonKit.git", .branch("master")),
-  ],
+  ] + conditionalPackageDependencies,
   targets: [
     .target(
       name: "CTensorFlow",
@@ -58,17 +77,18 @@ let package = Package(
         "CTensorFlow",
         "CX10Modules",
         .product(name: "Numerics", package: "swift-numerics"),
-      ],
+      ] + conditionalTargetDependencies,
       swiftSettings: [
         .define("DEFAULT_BACKEND_EAGER"),
-      ]),
+      ] + conditionalSwiftSettings),
     .target(
       name: "x10_optimizers_tensor_visitor_plan",
       dependencies: ["TensorFlow"],
       path: "Sources/x10",
       sources: [
         "swift_bindings/optimizers/TensorVisitorPlan.swift",
-      ]),
+      ],
+      swiftSettings: conditionalSwiftSettings),
     .target(
       name: "x10_optimizers_optimizer",
       dependencies: [
@@ -79,23 +99,29 @@ let package = Package(
       sources: [
         "swift_bindings/optimizers/Optimizer.swift",
         "swift_bindings/optimizers/Optimizers.swift",
-      ]),
-    // .target(
-    //   name: "x10_training_loop",
-    //   dependencies: ["TensorFlow"],
-    //   path: "Sources/x10",
-    //   sources: [
-    //     "swift_bindings/training_loop.swift",
-    //   ]),
+      ],
+      swiftSettings: conditionalSwiftSettings),
+     .target(
+       name: "x10_training_loop",
+       dependencies: ["TensorFlow"],
+       path: "Sources/x10",
+       sources: [
+         "swift_bindings/training_loop.swift",
+       ],
+       swiftSettings: conditionalSwiftSettings),
     .target(
       name: "Experimental",
-      dependencies: [],
-      path: "Sources/third_party/Experimental"),
+      dependencies: conditionalTargetDependencies,
+      path: "Sources/third_party/Experimental",
+      swiftSettings: conditionalSwiftSettings),
     .testTarget(
       name: "ExperimentalTests",
-      dependencies: ["Experimental"]),
+      dependencies: ["Experimental"],
+      swiftSettings: conditionalSwiftSettings),
     .testTarget(
       name: "TensorFlowTests",
-      dependencies: ["TensorFlow"]),
+      dependencies: ["TensorFlow"],
+      exclude: [],
+      swiftSettings: conditionalSwiftSettings),
   ]
 )
diff --git a/Sources/TensorFlow/BackwardsCompatibility.swift b/Sources/TensorFlow/BackwardsCompatibility.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 //===------------------------------------------------------------------------------------------===//
 // Losses

diff --git a/Sources/TensorFlow/Core/DataTypes.swift b/Sources/TensorFlow/Core/DataTypes.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 import CTensorFlow
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 import Numerics

diff --git a/Sources/TensorFlow/Core/DifferentialOperators.swift b/Sources/TensorFlow/Core/DifferentialOperators.swift
@@ -16,7 +16,11 @@
 // Free-function-style differential operators
 // ===------------------------------------------------------------------------------------------===//
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 // Value with gradient
 

diff --git a/Sources/TensorFlow/Core/ElementaryFunctions.swift b/Sources/TensorFlow/Core/ElementaryFunctions.swift
@@ -15,15 +15,21 @@
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
 import Numerics
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 extension ElementaryFunctions {
   internal static func visitChildren(
     _ body: (PartialKeyPath<Self>, ElementaryFunctionsVisit.Type) -> Void
   ) {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     if !_forEachFieldWithKeyPath(
       of: Self.self,

diff --git a/Sources/TensorFlow/Core/EuclideanDifferentiable.swift b/Sources/TensorFlow/Core/EuclideanDifferentiable.swift
@@ -12,11 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 func listFields<Root>(of type: Root.Type) -> [(String, PartialKeyPath<Root>)] {
   guard #available(macOS 9999, *) else {

diff --git a/Sources/TensorFlow/Core/KeyPathIterable.swift b/Sources/TensorFlow/Core/KeyPathIterable.swift
@@ -18,11 +18,19 @@
 // KeyPathIterable
 //===----------------------------------------------------------------------===//
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 /// An implementation detail of `KeyPathIterable`; do not use this protocol
 /// directly.
@@ -43,9 +51,11 @@ public protocol KeyPathIterable: _KeyPathIterableBase {
 
 public extension KeyPathIterable {
   var allKeyPaths: [PartialKeyPath<Self>] {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     var out = [PartialKeyPath<Self>]()
     _forEachFieldWithKeyPath(of: Self.self, options: .ignoreUnknown) { name, kp in

diff --git a/Sources/TensorFlow/Core/MixedPrecision.swift b/Sources/TensorFlow/Core/MixedPrecision.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 @_implementationOnly import x10_xla_tensor_wrapper
 

diff --git a/Sources/TensorFlow/Core/PointwiseMultiplicative.swift b/Sources/TensorFlow/Core/PointwiseMultiplicative.swift
@@ -12,11 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 infix operator .*: MultiplicationPrecedence
 infix operator .*=: AssignmentPrecedence
@@ -115,9 +123,11 @@ extension PointwiseMultiplicative {
   internal static func visitChildren(
     _ body: (PartialKeyPath<Self>, _PointwiseMultiplicative.Type) -> Void
   ) {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     if !_forEachFieldWithKeyPath(
       of: Self.self,

diff --git a/Sources/TensorFlow/Core/Tensor.swift b/Sources/TensorFlow/Core/Tensor.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 import CTensorFlow
 
 infix operator .==: ComparisonPrecedence

diff --git a/Sources/TensorFlow/Core/TensorGroup.swift b/Sources/TensorFlow/Core/TensorGroup.swift
@@ -341,12 +341,18 @@ extension Array: TensorArrayProtocol where Element: TensorGroup {
 
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 func reflectionInit<T>(type: T.Type, body: (inout T, PartialKeyPath<T>) -> Void) -> T {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
   guard #available(macOS 9999, *) else {
     fatalError("\(#function) is unavailable")
   }
+#endif
 
   let x = UnsafeMutablePointer<T>.allocate(capacity: 1)
   defer { x.deallocate() }
@@ -361,9 +367,11 @@ func reflectionInit<T>(type: T.Type, body: (inout T, PartialKeyPath<T>) -> Void)
 
 extension TensorGroup {
   public static var _typeList: [TensorDataType] {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     var out = [TensorDataType]()
     if !(_forEachFieldWithKeyPath(of: Self.self) { name, kp in
@@ -400,9 +408,11 @@ extension TensorGroup {
   }
 
   public func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?) {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     var i = 0
     if !_forEachFieldWithKeyPath(of: Self.self, body: { name, kp in

diff --git a/Sources/TensorFlow/Core/VectorProtocol.swift b/Sources/TensorFlow/Core/VectorProtocol.swift
@@ -12,11 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 
+#if canImport(ReflectionMirror)
+@_spi(Reflection) import ReflectionMirror
+#else
 @_spi(Reflection) import Swift
+#endif
 
 /// Implementation detail for reflection.
 ///
@@ -39,9 +47,11 @@ extension VectorProtocol {
   internal static func visitChildren(
     _ body: (PartialKeyPath<Self>, _VectorProtocol.Type) -> Void
   ) {
+#if !TENSORFLOW_USE_RELEASE_TOOLCHAIN
     guard #available(macOS 9999, *) else {
       fatalError("\(#function) is unavailable")
     }
+#endif
 
     if !_forEachFieldWithKeyPath(
       of: Self.self,

diff --git a/Sources/TensorFlow/Initializers.swift b/Sources/TensorFlow/Initializers.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 
 extension Tensor {
   /// Creates a tensor with the specified shape and a single, repeated scalar value.

diff --git a/Sources/TensorFlow/Layer.swift b/Sources/TensorFlow/Layer.swift
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#if canImport(Differentiation)
+import Differentiation
+#else
 import _Differentiation
+#endif
 import Foundation
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
 import Numerics
@@ -239,9 +243,15 @@ extension Layer {
     -> (output: Output, backpropagator: Backpropagator)
   {
 #if TENSORFLOW_USE_STANDARD_TOOLCHAIN
+    #if canImport(Differentiation)
+    let (out, pullback) = Differentiation.valueWithPullback(at: self, input) { layer, input in
+      return layer(input)
+    }
+    #else
     let (out, pullback) = _Differentiation.valueWithPullback(at: self, input) { layer, input in
       return layer(input)
     }
+    #endif
 #else
     let (out, pullback) = Swift.valueWithPullback(at: self, input) { layer, input in
       return layer(input)