[TF] WIP: Removing all instances of #tfop. (#24425)

* Removed some instances of '#tfop'.

* Removed all instances of '#tfop'.

* Fixes to the build scripts for the new swift-bindings.

* Disabled a couple of GPE tests.

* Added back a couple of test helpers.

* Minor tweak.

* Minor tweaks.

* Removed all uses of '.toHost' and '.toAccelerator'.

* Added back some tests.

* Bug fixes.

* Updated the swift-bindings dependency.

Author: eaplatanios

README.md

@@ -151,8 +151,8 @@ Below is more information about TensorFlow-related build arguments.
     * Default value: None.
 * `tensorflow-swift-apis`: A path to the [tensorflow/swift-apis](https://github.com/tensorflow/swift-apis) deep learning library repository.
     * Default value: `tensorflow-swift-apis` if the [tensorflow/swift-apis](https://github.com/tensorflow/swift-apis) repository is cloned. Otherwise, none.
-* `tensorflow-swift-bindings`: A generated TensorFlow Swift bindings file (`RawOpsGenerated.swift`) obtained from [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings).
-    * Default value: `tensorflow-swift-bindings/RawOpsGenerated.swift` if the [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings) repository is cloned. Otherwise, none.
+* `tensorflow-swift-bindings`: A path to the [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings) repository.
+    * Default value: `tensorflow-swift-bindings` if the [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings) repository is cloned. Otherwise, none.
 
 ### Build systems
 

stdlib/public/TensorFlow/ArrayOps.swift

@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains some Array ops that cannot be properly handled by #tfop.
+// This file contains some custom raw ops.
 //
 // TODO: These should be deleted once we can properly generate raw ops for these.
 //
@@ -109,91 +109,4 @@ public extension Raw {
     }
     return out
   }
-
-  /// Splits a tensor into `numSplit` tensors along one dimension.
-  ///
-  /// - Parameters:
-  ///   - splitDim: 0-D.  The dimension along which to split.  Must be in the range
-  ///     `[-rank(value), rank(value))`.
-  ///   - value: The tensor to split.
-  ///   - numSplit: The number of splits to create.
-  ///
-  /// - Returns: Tensors whose shape matches that of `value`
-  ///   except along `axis`, where their sizes are
-  ///   `value.shape[axis] / numSplit`.
-  @inlinable @inline(__always)
-  static func split<T: TensorFlowScalar>(
-    splitDim: Tensor<Int32>,
-    value: Tensor<T>,
-    numSplit: Int64
-  ) -> [Tensor<T>] {
-    let s: CTFStatus = TF_NewStatus()
-    defer { TF_DeleteStatus(s) }
-    let op: CTFEOp = TFE_NewOp(_ExecutionContext.global.eagerContext, "Split", s)
-    defer { TFE_DeleteOp(op) }
-    let _ = _TFCOpAddInputFromTensorGroup(op, splitDim, s)
-    let _ = _TFCOpAddInputFromTensorGroup(op, value, s)
-    TFE_OpSetAttrInt(op, "num_split", numSplit)
-    TFE_OpSetAttrType(op, "T", T.tensorFlowDataType._cDataType)
-    var count: Int32 = Int32(numSplit)
-    let buffer: UnsafeMutablePointer<CTensorHandle> =
-      UnsafeMutablePointer.allocate(capacity: Int(count))
-    defer { buffer.deallocate() }
-    _TFCEagerExecute(op, UnsafeMutablePointer<CTensorHandle?>(buffer), &count, s)
-    checkOk(s)
-
-    var out: [Tensor<T>] = []
-    var cursor = buffer
-    for _ in 0..<numSplit {
-      out.append(Tensor<T>(handle: TensorHandle(_owning: cursor.pointee)))
-      cursor = cursor.advanced(by: 1)
-    }
-    return out
-  }
-
-  /// Splits a tensor into `numSplit` tensors along one dimension.
-  ///
-  /// - Parameters:
-  ///   - value: The tensor to split.
-  ///   - sizeSplits: list containing the sizes of each output tensor along the split
-  ///     dimension. Must sum to the dimension of value along split_dim.
-  ///     Can contain one -1 indicating that dimension is to be inferred.
-  ///   - splitDim: 0-D.  The dimension along which to split.  Must be in the range
-  ///     `[-rank(value), rank(value))`.
-  ///
-  /// - Returns: Tensors whose shape matches that of `value`
-  ///   except along `axis`, where their sizes are
-  ///   `size_splits[i]`.
-  @inlinable @inline(__always)
-  static func splitV<T: TensorFlowScalar, Tlen: BinaryInteger & TensorFlowScalar>(
-    value: Tensor<T>,
-    sizeSplits: Tensor<Tlen>,
-    splitDim: Tensor<Int32>,
-    numSplit: Int64
-  ) -> [Tensor<T>] {
-    let s: CTFStatus = TF_NewStatus()
-    defer { TF_DeleteStatus(s) }
-    let op: CTFEOp = TFE_NewOp(_ExecutionContext.global.eagerContext, "SplitV", s)
-    defer { TFE_DeleteOp(op) }
-    let _ = _TFCOpAddInputFromTensorGroup(op, value, s)
-    let _ = _TFCOpAddInputFromTensorGroup(op, sizeSplits, s)
-    let _ = _TFCOpAddInputFromTensorGroup(op, splitDim, s)
-    TFE_OpSetAttrInt(op, "num_split", numSplit)
-    TFE_OpSetAttrType(op, "T", T.tensorFlowDataType._cDataType)
-    TFE_OpSetAttrType(op, "Tlen", Tlen.tensorFlowDataType._cDataType)
-    var count: Int32 = Int32(numSplit)
-    let buffer: UnsafeMutablePointer<CTensorHandle> =
-      UnsafeMutablePointer.allocate(capacity: Int(count))
-    defer { buffer.deallocate() }
-    _TFCEagerExecute(op, UnsafeMutablePointer<CTensorHandle?>(buffer), &count, s)
-    checkOk(s)
-
-    var out: [Tensor<T>] = []
-    var cursor = buffer
-    for _ in 0..<numSplit {
-      out.append(Tensor<T>(handle: TensorHandle(_owning: cursor.pointee)))
-      cursor = cursor.advanced(by: 1)
-    }
-    return out
-  }
 }

stdlib/public/TensorFlow/CMakeLists.txt

@@ -56,7 +56,9 @@ set(SOURCES
 
 # Copy TensorFlow bindings file, if it exists.
 if (TENSORFLOW_SWIFT_BINDINGS)
-  list(APPEND SOURCES "${TENSORFLOW_SWIFT_BINDINGS}")
+  file(GLOB_RECURSE TENSORFLOW_SWIFT_BINDINGS_SOURCES
+    "${TENSORFLOW_SWIFT_BINDINGS}/*.swift")
+  list(APPEND SOURCES "${TENSORFLOW_SWIFT_BINDINGS_SOURCES}")
 endif()
 
 # Copy TensorFlow high-level API sources, if they exist.

stdlib/public/TensorFlow/Dataset.swift

@@ -58,26 +58,21 @@ public extension Dataset {
   @inlinable @inline(__always)
   init(randomSeed: Int64) {
     let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
-    self.init(
-      _handle: #tfop("RandomDataset", seed1, seed2,
-                     output_types$dtype: Element._typeList,
-                     output_shapes: Element._unknownShapeList)
-    )
+    self.init(_handle: Raw.experimentalRandomDataset(
+      seed: seed1,
+      seed2: seed2,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
   }
 }
 
 public extension Dataset {
   /// Creates a dataset from a batch of elements as a tensor.
   @inlinable @inline(__always)
   init(elements: Element) {
-    // A dataset creation op only runs on TF CPU.
-    self.init(
-      _handle: #tfop(
-        "TensorSliceDataset", [elements],
-        Toutput_types$dtype: Element._typeList,
-        output_shapes: Element._unknownShapeList
-      )
-    )
+    self.init(_handle: Raw.tensorSliceDataset(
+      components: [elements],
+      outputShapes: Element._unknownShapeList))
   }
 }
 
@@ -87,10 +82,10 @@ extension Dataset : Sequence {
   /// Returns an iterator over the elements of this dataset.
   @inlinable @inline(__always)
   public func makeIterator() -> DatasetIterator<Element> {
-    let resource: ResourceHandle =
-      #tfop("AnonymousIterator", output_types$dtype: Element._typeList,
-            output_shapes: Element._unknownShapeList)
-    #tfop("MakeIterator", _handle, resource) as Void
+    let resource = Raw.anonymousIterator(
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList)
+    Raw.makeIterator(dataset: _handle, iterator: resource)
     return DatasetIterator(_handle: resource)
   }
 }
@@ -102,44 +97,43 @@ public extension Dataset {
   func map<ResultElement : TensorGroup>(
     _ transform: (Element) -> ResultElement
   ) -> Dataset<ResultElement> {
-    return Dataset<ResultElement>(
-      _handle: #tfop(
-        "MapDataset", _handle, [Tensor<Int32>(0)],
-        f$func: _tffunc(transform),
-        Targuments$dtype: [Int32.tensorFlowDataType],
-        output_types$dtype: ResultElement._typeList,
-        output_shapes: ResultElement._unknownShapeList
-      )
-    )
+    return Dataset<ResultElement>(_handle: Raw.mapDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      f: transform,
+      outputTypes: ResultElement._typeList,
+      outputShapes: ResultElement._unknownShapeList,
+      useInterOpParallelism: true,
+      preserveCardinality: false))
   }
 
   @inlinable @inline(__always)
-  func map<ResultElement : TensorGroup>(parallelCallCount: Int,
-    _ transform: (Element) -> ResultElement) -> Dataset<ResultElement> {
-    return Dataset<ResultElement>(
-      _handle: #tfop("ParallelMapDataset", _handle, [Tensor<Int32>(0)],
-        [Tensor<Int32>(Int32(parallelCallCount))],
-        f$func: _tffunc(transform),
-        Targuments$dtype: [Int32.tensorFlowDataType],
-        output_types$dtype: ResultElement._typeList,
-        output_shapes: ResultElement._unknownShapeList
-      )
-    )
+  func map<ResultElement : TensorGroup>(
+    parallelCallCount: Int,
+    _ transform: (Element) -> ResultElement
+  ) -> Dataset<ResultElement> {
+    return Dataset<ResultElement>(_handle: Raw.parallelMapDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      numParallelCalls: Tensor<Int32>(Int32(parallelCallCount)),
+      f: transform,
+      outputTypes: ResultElement._typeList,
+      outputShapes: ResultElement._unknownShapeList,
+      useInterOpParallelism: true,
+      sloppy: false,
+      preserveCardinality: false))
   }
 
   @inlinable @inline(__always)
   func filter(
     _ isIncluded: (Element) -> Tensor<Bool>
   ) -> Dataset {
-    return Dataset(
-      _handle: #tfop(
-        "FilterDataset", _handle, [Tensor<Int32>(0)],
-        predicate$func: _tffunc(isIncluded),
-        Targuments$dtype: [Int32.tensorFlowDataType],
-        output_types$dtype: Element._typeList,
-        output_shapes: Element._unknownShapeList
-      )
-    )
+    return Dataset(_handle: Raw.filterDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      predicate: isIncluded,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
   }
 }
 
@@ -149,24 +143,22 @@ public extension Dataset {
     sampleCount: Int, randomSeed: Int64
   ) -> Dataset {
     let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
-    return Dataset(
-      _handle: #tfop(
-        "ShuffleDataset", _handle, Tensor(Int64(sampleCount)), seed1, seed2,
-        output_types$dtype: Element._typeList,
-        output_shapes: Element._unknownShapeList
-      )
-    )
+    return Dataset(_handle: Raw.shuffleDataset(
+      inputDataset: _handle,
+      bufferSize: Tensor(Int64(sampleCount)),
+      seed: seed1,
+      seed2: seed2,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
   }
 
   @inlinable @inline(__always)
   func batched(_ batchSize: Int) -> Dataset {
-    return Dataset(
-      _handle: #tfop(
-        "BatchDataset", _handle, Tensor(Int64(batchSize)),
-        output_types$dtype: Element._typeList,
-        output_shapes: Element._unknownShapeList
-      )
-    )
+    return Dataset(_handle: Raw.batchDataset(
+      inputDataset: _handle,
+      batchSize: Tensor(Int64(batchSize)),
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
   }
 }
 
@@ -186,16 +178,16 @@ extension DatasetIterator : IteratorProtocol {
   /// exists.
   @inlinable @inline(__always)
   public mutating func next() -> Element? {
-    let optional: VariantHandle =
-      #tfop("IteratorGetNextAsOptional", _handle,
-            output_types$dtype: Element._typeList,
-            output_shapes: Element._unknownShapeList)
-    guard _TFGetScalarOrDie(#tfop("OptionalHasValue", optional)) else {
+    let optional = Raw.iteratorGetNextAsOptional(
+      iterator: _handle,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList)
+    guard Raw.optionalHasValue(optional: optional).scalarized() else {
       return nil
     }
-    return #tfop("OptionalGetValue", optional,
-                 output_types$dtype: Element._typeList,
-                 output_shapes: Element._unknownShapeList) as Element
+    return Raw.optionalGetValue(
+      optional: optional,
+      outputShapes: Element._unknownShapeList)
   }
 }
 
@@ -217,9 +209,9 @@ public struct Zip2TensorGroup<T : TensorGroup, U : TensorGroup> : TensorGroup {
 public func zip<T : TensorGroup, U : TensorGroup>(
   _ dataset1: Dataset<T>, _ dataset2: Dataset<U>
 ) -> Dataset<Zip2TensorGroup<T, U>> {
-  let handle: VariantHandle = #tfop(
-     "ZipDataset", Zip2TensorGroup(dataset1._handle, dataset2._handle),
-     output_types$dtype: Zip2TensorGroup<T, U>._typeList,
-     output_shapes: Zip2TensorGroup<T, U>._unknownShapeList)
+  let handle = Raw.zipDataset(
+    inputDatasets: [dataset1._handle, dataset2._handle],
+    outputTypes: Zip2TensorGroup<T, U>._typeList,
+    outputShapes: Zip2TensorGroup<T, U>._unknownShapeList)
   return Dataset(_handle: handle)
 }

stdlib/public/TensorFlow/Ops.swift

@@ -25,18 +25,6 @@
 // we also define some helper function wrappers, e.g. to make things symmetric
 // and generally feel nice to use.
 //
-// The ops themselves are defined by the primitive #tfop(...) syntax, here are
-// some examples:
-//     result = #tfop("Add", lhs, rhs)
-//     result = #tfop("Const", dtype: Float.self, value$tensor: 4.0)
-//
-// The first parameter to this syntax is the TensorFlow op name as a string.
-// After that, the inputs are specified, and then attributes are specified
-// with their name as the keyword argument.
-//
-// Inputs and outputs must be of TensorHandle, ResourceHandle, or VariantHandle
-// type.  These are magic types known to the compiler.
-//
 
 infix operator ++ : AdditionPrecedence
 
@@ -700,17 +688,15 @@ internal extension Tensor where Scalar : TensorFlowFloatingPoint {
   @inlinable @inline(__always)
   func _vjpConcatenated(with other: Tensor, alongAxis axis: Int)
     -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
-    let idx = axis < 0 ? axis + rank : axis
-    let splits = Tensor<Int32>([shapeTensor[idx], other.shapeTensor[idx]])
+    let posAxis = axis < 0 ? axis + rank: axis
+    let splits = Tensor<Int32>([shapeTensor[posAxis], other.shapeTensor[posAxis]])
     return (concatenated(with: other, alongAxis: axis), { result in
-      let ret: (TensorHandle<Scalar>, TensorHandle<Scalar>) = #tfop("SplitV",
-        result,
-        splits,
-        Tensor<Int32>(Int32(axis)),
-        num_split: Int64(2),
-        T$dtype: Scalar.tensorFlowDataType,
-        Tlen$dtype: Int32.tensorFlowDataType)
-      return (Tensor(handle: ret.0), Tensor(handle: ret.1))
+      let gradients = Raw.splitV(
+        value: result,
+        sizeSplits: splits,
+        splitDim: Tensor<Int32>(Int32(axis)),
+        numSplit: Int64(splits.shape[0]))
+      return (gradients[0], gradients[1])
     })
   }
 }

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)
   }
 }