Changes to support the new swift-bindings.

Author: eaplatanios

stdlib/public/TensorFlow/StringOps.swift

@@ -23,7 +23,6 @@ public extension StringTensor {
   /// - Note: `elementsEqual` supports broadcasting.
   @inlinable @inline(__always)
   func elementsEqual(_ other: StringTensor) -> Tensor<Bool> {
-    return #tfop("Equal", self.handle, other.handle,
-                 T$dtype: String.tensorFlowDataType)
+    return Raw.equal(self, other)
   }
 }

stdlib/public/TensorFlow/Tensor.swift

@@ -416,12 +416,7 @@ extension _TensorElementLiteral : ExpressibleByArrayLiteral {
   public typealias ArrayLiteralElement = _TensorElementLiteral<Scalar>
   @inlinable @inline(__always)
   public init(arrayLiteral elements: _TensorElementLiteral<Scalar>...) {
-    // Attr T (non-optional in the op definition) need not be specified when we
-    // run the op as part of a graph function, but need to be specified when we
-    // run it via eager C API.
-    let handle: TensorHandle<Scalar> = #tfop("Pack", elements,
-                                             T$dtype: Scalar.tensorFlowDataType)
-    tensor = Tensor(handle: handle)
+    tensor = Raw.pack(elements.map { $0.tensor })
   }
 }
 
@@ -436,8 +431,7 @@ extension Tensor : ExpressibleByArrayLiteral {
   internal init(
     _tensorElementLiterals elements: [_TensorElementLiteral<Scalar>]
   ) {
-    self.init(handle: #tfop("Pack", elements,
-                            T$dtype: Scalar.tensorFlowDataType))
+    self = Raw.pack(elements.map { $0.tensor })
   }
 
   /// Creates a tensor initialized with the given elements.

stdlib/public/TensorFlow/TensorGroup.swift

@@ -29,6 +29,8 @@ public protocol TensorArrayProtocol {
   func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?)
 
   var _tensorHandleCount: Int32 { get }
+
+  init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int)
 }
 
 /// A protocol representing types that can be mapped to and from
@@ -67,6 +69,11 @@ public extension TensorGroup {
   static var _unknownShapeList: [TensorShape?] {
     return Array(repeating: nil, count: _typeList.count)
   }
+
+  init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int) {
+    precondition(count == Self._typeList.count)
+    self.init(_owning: tensorHandles)
+  }
 }
 
 //===----------------------------------------------------------------------===//
@@ -199,7 +206,7 @@ extension StringTensor : TensorGroup {
   }
 }
 
-extension Array : TensorArrayProtocol where Element : TensorArrayProtocol {
+extension Array : TensorArrayProtocol where Element : TensorGroup {
   public func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?) {
     var ptr = address
     for elem in self {
@@ -213,4 +220,11 @@ extension Array : TensorArrayProtocol where Element : TensorArrayProtocol {
     for elem in self { count += elem._tensorHandleCount }
     return count
   }
+
+  public init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int) {
+    let size = count / Int(Element._tensorHandleCount)
+    self = Array((0..<size).map { Element.init(
+      _owning: tensorHandles?.advanced(by: $0 * Int(Element._tensorHandleCount)))
+    })
+  }
 }

test/TensorFlowRuntime/tracer.swift

@@ -168,13 +168,25 @@ TracerTests.testAllBackends("Advanced") {
     var model: Model = [Tensor<Float>(1.0), Tensor<Float>(2.0)]
     var optimizer: Optimizer = [Tensor<Float>(1.0), Tensor<Float>(2.0)]
 
+    public init() {}
+
+    public init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int) {
+      self.model = [
+        Tensor<Float>(_owning: tensorHandles),
+        Tensor<Float>(_owning: tensorHandles?.advanced(by: 1))]
+      self.optimizer = [
+        Tensor<Float>(_owning: tensorHandles?.advanced(by: 2)),
+        Tensor<Float>(_owning: tensorHandles?.advanced(by: 3))]
+    }
+
     public func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?) {
       print("Calling State._unpackTensorHandles().")
       var ptr = address
       model._unpackTensorHandles(into: ptr)
       ptr = ptr!.advanced(by: Int(model._tensorHandleCount))
       optimizer._unpackTensorHandles(into: ptr)
     }
+
     public var _tensorHandleCount: Int32 {
       return model._tensorHandleCount + optimizer._tensorHandleCount
     }

utils/build-script-impl

@@ -280,7 +280,7 @@ KNOWN_SETTINGS=(
     tensorflow-host-include-dir        ""        "Path to host TensorFlow headers"
     tensorflow-target-include-dir      ""        "Path to target Tensorflow headers"
     tensorflow-target-lib-dir          ""        "Path to target TensorFlow libraries"
-    tensorflow-swift-bindings          ""        "Path to TensorFlow Swift bindings file"
+    tensorflow-swift-bindings          ""        "Path to TensorFlow Swift bindings repository"
     tensorflow-swift-apis              ""        "Path to TensorFlow deep learning library repository"
 )
 
@@ -2476,7 +2476,7 @@ for host in "${ALL_HOSTS[@]}"; do
 
                     # Handle TensorFlow Swift bindings file.
                     if [[ ! "${TENSORFLOW_SWIFT_BINDINGS}" && -d "${TENSORFLOW_SWIFT_BINDINGS_DIR}" ]] ; then
-                        TENSORFLOW_SWIFT_BINDINGS="${TENSORFLOW_SWIFT_BINDINGS_DIR}/RawOpsGenerated.swift"
+                        TENSORFLOW_SWIFT_BINDINGS="${TENSORFLOW_SWIFT_BINDINGS_DIR}"
                     fi
                     if [[ "${TENSORFLOW_SWIFT_BINDINGS}" ]] ; then
                         cmake_options=(

utils/build_swift/driver_arguments.py

@@ -974,8 +974,7 @@ def create_argument_parser():
                 'Used for linking Swift programs.')
     option('--tensorflow-swift-bindings', store_path,
            default=None,
-           help='Path to a TensorFlow Swift bindings file '
-                '(RawOpsGenerated.swift).')
+           help='Path to a TensorFlow Swift bindings repository.')
     option('--tensorflow-swift-apis', store_path,
            default=None,
            help='Path to a TensorFlow deep learning library repository.')