Merge branch 'master' into complex-numbers2

Author: bartchr808

.gitignore

@@ -5,6 +5,8 @@ xcuserdata
 DerivedData/
 *.xcodeproj
 *~
+*.vscode
+*.idea
 
 ### MacOS ###
 .DS_Store

Dockerfile

@@ -3,7 +3,7 @@
 FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04
 
 # Allows the caller to specify the toolchain to use.
-ARG swift_tf_url=https://storage.googleapis.com/s4tf-kokoro-artifact-testing/latest/swift-tensorflow-DEVELOPMENT-cuda10.0-cudnn7-ubuntu18.04.tar.gz
+ARG swift_tf_url=https://storage.googleapis.com/s4tf-kokoro-artifact-testing/latest/swift-tensorflow-DEVELOPMENT-cuda10.0-cudnn7-test-ubuntu18.04.tar.gz
 
 # Install Swift deps.
 ENV DEBIAN_FRONTEND=noninteractive

Sources/TensorFlow/Core/DataTypes.swift

@@ -14,14 +14,18 @@
 
 import CTensorFlow
 
-public extension TensorDataType {
-    var _cDataType: TF_DataType {
-        return TF_DataType(rawValue: _internalStorageType)
+/// A TensorFlow dynamic type value that can be created from types that conform to
+/// `TensorFlowScalar`.
+// This simply wraps a `TF_DataType` and allows user code to handle
+// `TF_DataType` without importing CTensorFlow, which pollutes the namespace
+// with TensorFlow C API declarations.
+public struct TensorDataType {
+    public var _cDataType: TF_DataType
+
+    @usableFromInline
+	internal init(_ cDataType: TF_DataType) {
+	    self._cDataType = cDataType
     }
-
-    init(_ cDataType: TF_DataType) {
-	    self.init(rawValue: cDataType.rawValue)
-	}
 }
 
 @usableFromInline

Sources/TensorFlow/Core/DifferentialOperators.swift

@@ -21,14 +21,15 @@ public extension Differentiable {
     func gradient<R: TensorFlowFloatingPoint>(
         in f: @differentiable (Self) -> Tensor<R>
     ) -> TangentVector {
-        return self.pullback(in: f)(Tensor<R>(1))
+        return self.valueWithGradient(in: f).1
     }
 
     @inlinable
     func valueWithGradient<R: TensorFlowFloatingPoint>(
         in f: @differentiable (Self) -> Tensor<R>
     ) -> (value: Tensor<R>, gradient: TangentVector) {
         let (y, pb) = self.valueWithPullback(in: f)
+        precondition(y.rank == 0)
         return (y, pb(Tensor<R>(1)))
     }
 
@@ -37,7 +38,7 @@ public extension Differentiable {
         at x: T,
         in f: @differentiable (Self, T) -> Tensor<R>
     ) -> (TangentVector, T.TangentVector) {
-        return self.pullback(at: x, in: f)(Tensor<R>(1))
+        return self.valueWithGradient(at: x, in: f).1
     }
 
     @inlinable
@@ -46,6 +47,7 @@ public extension Differentiable {
         in f: @differentiable (Self, T) -> Tensor<R>
     ) -> (value: Tensor<R>, gradient: (TangentVector, T.TangentVector)) {
         let (y, pb) = self.valueWithPullback(at: x, in: f)
+        precondition(y.rank == 0)
         return (y, pb(Tensor<R>(1)))
     }
 }
@@ -63,6 +65,7 @@ public func valueWithGradient<T, R>(
 ) -> (value: Tensor<R>, gradient: T.TangentVector)
 where T: Differentiable, R: TensorFlowFloatingPoint {
     let (y, pullback) = valueWithPullback(at: x, in: f)
+    precondition(y.rank == 0)
     return (y, pullback(Tensor<R>(1)))
 }
 
@@ -74,6 +77,7 @@ public func valueWithGradient<T, U, R>(
 ) -> (value: Tensor<R>, gradient: (T.TangentVector, U.TangentVector))
     where T: Differentiable, U: Differentiable, R: TensorFlowFloatingPoint {
     let (y, pullback) = valueWithPullback(at: x, y, in: f)
+    precondition(y.rank == 0)
     return (y, pullback(Tensor<R>(1)))
 }
 
@@ -86,6 +90,7 @@ public func valueWithGradient<T, U, R>(
 // ) -> (value: Tensor<R>, gradient: (T.TangentVector, U.TangentVector, V.TangentVector))
 //   where T: Differentiable, U: Differentiable, V: Differentiable, R: TensorFlowFloatingPoint {
 //   let (y, pullback) = valueWithPullback(at: x, y, z, in: f)
+//   precondition(y.rank == 0)
 //   return (y, pullback(Tensor<R>(1)))
 // }
 
@@ -124,7 +129,7 @@ public func gradient<T, R>(
     at x: T,
     in f: @differentiable (T) -> Tensor<R>
 ) -> T.TangentVector where T: Differentiable, R: TensorFlowFloatingPoint {
-    return pullback(at: x, in: f)(Tensor<R>(1))
+    return valueWithGradient(at: x, in: f).1
 }
 
 @inlinable
@@ -134,7 +139,7 @@ public func gradient<T, U, R>(
     in f: @differentiable (T, U) -> Tensor<R>
 ) -> (T.TangentVector, U.TangentVector)
     where T: Differentiable, U: Differentiable, R: TensorFlowFloatingPoint {
-    return pullback(at: x, y, in: f)(Tensor<R>(1))
+    return valueWithGradient(at: x, y, in: f).1
 }
 
 // @inlinable
@@ -145,7 +150,7 @@ public func gradient<T, U, R>(
 //     in f: @differentiable (T, U, V) -> Tensor<R>
 // ) -> (T.TangentVector, U.TangentVector, V.TangentVector)
 //     where T: Differentiable, U: Differentiable, V: Differentiable, R: TensorFlowFloatingPoint {
-//     return pullback(at: x, y, z, in: f)(Tensor<R>(1))
+//     return valueWithGradient(at: x, y, z, in: f).1
 // }
 
 // Gradient (curried)

Sources/TensorFlow/Core/Runtime.swift

@@ -551,7 +551,7 @@ public final class _ExecutionContext {
     @usableFromInline let eagerContext: CTFEContext
 
     /// The status for checking TensorFlow errors.
-    private let status: CTFStatus = TF_NewStatus()
+    @usableFromInline let status: CTFStatus = TF_NewStatus()
 
     /// The mutex for preventing potential concurrent access.
     private var mutex: pthread_mutex_t = pthread_mutex_t()
@@ -569,6 +569,12 @@ public final class _ExecutionContext {
         // Initialize the TF runtime exactly once. Only affects local execution
         // (when _RuntimeConfig.tensorFlowServer is set to "").
         if !_RuntimeConfig.tensorFlowRuntimeInitialized {
+            // Install a signal handler to ensure we exit when interrupted.
+            signal(SIGINT) { _ in
+                print("Caught interrupt signal, exiting...")
+                exit(1)
+            }
+            
             var args = ["dummyProgramName"]
             if _RuntimeConfig.printsDebugLog {
                 args.append("--alsologtostderr")
@@ -588,24 +594,19 @@ public final class _ExecutionContext {
 
             // Calculate the addresses of all the strings within our single buffer, and then call
             // TF_InitMain.
-            flattenedStringBytes.withUnsafeMutableBufferPointer { flattenedStringBytesBuffer in
+            flattenedStringBytes.withUnsafeMutableBufferPointer { buffer in
                 var stringAddrs: [UnsafeMutablePointer<Int8>?] = []
-                var currentStringAddr = flattenedStringBytesBuffer.baseAddress
+                var currentStringAddr = buffer.baseAddress
                     .map(UnsafeMutablePointer.init)
                 for length in lengths {
                     stringAddrs.append(currentStringAddr)
                     currentStringAddr = currentStringAddr?.advanced(by: length)
                 }
 
                 stringAddrs.withUnsafeMutableBufferPointer { stringAddrsBuffer in
-                    var cArgs = [stringAddrsBuffer.baseAddress.map(UnsafeMutablePointer.init)]
-                    var cArgsCount = [Int32(args.count)]
-
-                    cArgs.withUnsafeMutableBufferPointer { cArgsBuffer in
-                        cArgsCount.withUnsafeMutableBufferPointer { cArgsCountBuffer in
-                            TF_InitMain(nil, cArgsCountBuffer.baseAddress, cArgsBuffer.baseAddress)
-                        }
-                    }
+                    var cArgsCount = Int32(args.count)
+                    var cArgs = stringAddrsBuffer.baseAddress.map(UnsafeMutablePointer.init)
+                    TF_InitMain(nil, &cArgsCount, &cArgs)
                 }
             }
             _RuntimeConfig.tensorFlowRuntimeInitialized = true

Sources/TensorFlow/Core/Tensor.swift

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import CTensorFlow
+
 infix operator .==: ComparisonPrecedence
 infix operator .!=: ComparisonPrecedence
 
@@ -50,7 +52,10 @@ public extension Tensor {
     var rank: Int {
         @_semantics("autodiff.nonvarying")
         get {
-            return Int(rankTensor.scalar!)
+            let status = _ExecutionContext.global.status
+            let rank = TFE_TensorHandleNumDims(handle._cTensorHandle, status)
+            checkOk(status)
+            return Int(rank)
         }
     }
 
@@ -59,15 +64,25 @@ public extension Tensor {
     var shape: TensorShape {
         @_semantics("autodiff.nonvarying")
         get {
-            return TensorShape(shapeTensor.scalars.map(Int.init))
+            let status = _ExecutionContext.global.status
+            let dims: [Int] = (0..<Int32(rank)).map { i in
+                let dim = TFE_TensorHandleDim(self.handle._cTensorHandle, i, status)
+                checkOk(status)
+                return Int(dim)
+            }
+            return TensorShape(dims)
         }
     }
 
     /// The number of scalars in the `Tensor`.
     @inlinable
     var scalarCount: Int {
+        @_semantics("autodiff.nonvarying")
         get {
-            return Int(scalarCountTensor.scalar!)
+            let status = _ExecutionContext.global.status
+            let size = TFE_TensorHandleNumElements(handle._cTensorHandle, status)
+            checkOk(status)
+            return Int(size)
         }
     }
 
@@ -511,20 +526,12 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
         lhs: Tensor,
         rhs: Tensor
     ) -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
-        return (lhs + rhs, { [
-            lhsShape = lhs.shape,
-            rhsShape = rhs.shape,
-            lhsShapeTensor = lhs.shapeTensor,
-            rhsShapeTensor = rhs.shapeTensor] v in
-            var lhsGrad = v
-            var rhsGrad = v
-            if lhsGrad.shape != lhsShape {
-                lhsGrad = lhsGrad.unbroadcasted(toShape: lhsShapeTensor)
-            }
-            if rhsGrad.shape != rhsShape {
-                rhsGrad = rhsGrad.unbroadcasted(toShape: rhsShapeTensor)
-            }
-            return (lhsGrad, rhsGrad)
+        return (lhs + rhs, { [lhsShape = lhs.shapeTensor, rhsShape = rhs.shapeTensor] v in
+            let lhsGrad = v
+            let rhsGrad = lhsGrad
+            let (lhsAxes, rhsAxes) = Raw.broadcastGradientArgs(s0: lhsShape, s1: rhsShape)
+            return (lhsGrad.sum(squeezingAxes: lhsAxes).reshaped(toShape: lhsShape),
+                    rhsGrad.sum(squeezingAxes: rhsAxes).reshaped(toShape: rhsShape))
         })
     }
 
@@ -533,20 +540,12 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
         lhs: Tensor,
         rhs: Tensor
     ) -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
-        return (lhs - rhs, { [
-            lhsShape = lhs.shape,
-            rhsShape = rhs.shape,
-            lhsShapeTensor = lhs.shapeTensor,
-            rhsShapeTensor = rhs.shapeTensor] v in
-            var lhsGrad = v
-            var rhsGrad = -v
-            if lhsGrad.shape != lhsShape {
-                lhsGrad = lhsGrad.unbroadcasted(toShape: lhsShapeTensor)
-            }
-            if rhsGrad.shape != rhsShape {
-                rhsGrad = rhsGrad.unbroadcasted(toShape: rhsShapeTensor)
-            }
-            return (lhsGrad, rhsGrad)
+        return (lhs - rhs, { [lhsShape = lhs.shapeTensor, rhsShape = rhs.shapeTensor] v in
+            let lhsGrad = v
+            let rhsGrad = -lhsGrad
+            let (lhsAxes, rhsAxes) = Raw.broadcastGradientArgs(s0: lhsShape, s1: rhsShape)
+            return (lhsGrad.sum(squeezingAxes: lhsAxes).reshaped(toShape: lhsShape),
+                    rhsGrad.sum(squeezingAxes: rhsAxes).reshaped(toShape: rhsShape))
         })
     }
 }

Sources/TensorFlow/Core/TensorGroup.swift

@@ -14,6 +14,45 @@
 
 import CTensorFlow
 
+/// A protocol representing types that can be mapped to `Array<CTensorHandle>`.
+///
+/// This protocol is defined separately from `TensorGroup` in order for the number of tensors to be 
+/// determined at runtime. For example, `[Tensor<Float>]` may have an unknown number of elements at 
+/// compile time.
+///
+/// This protocol can be derived automatically for structs whose stored properties all conform to 
+/// the `TensorGroup` protocol. It cannot be derived automatically for structs whose properties all 
+/// conform to `TensorArrayProtocol` due to the constructor requirement (i.e., in such cases it 
+/// would be impossible to know how to break down `count` among the stored properties).
+public protocol TensorArrayProtocol {
+    /// Writes the tensor handles to `address`, which must be allocated with enough capacity to hold 
+    /// `_tensorHandleCount` handles. The tensor handles written to `address` are borrowed: this 
+    /// container still owns them.
+    func _unpackTensorHandles(into address: UnsafeMutablePointer<CTensorHandle>?)
+
+    var _tensorHandleCount: Int32 { get }
+    var _typeList: [TensorDataType] { get }
+
+    init(_owning tensorHandles: UnsafePointer<CTensorHandle>?, count: Int)
+}
+
+/// A protocol representing types that can be mapped to and from `Array<CTensorHandle>`.
+///
+/// When a `TensorGroup` is used as an argument to a tensor operation, it is passed as an argument
+/// list whose elements are the tensor fields of the type.
+///
+/// When a `TensorGroup` is returned as a result of a tensor operation, it is initialized with its 
+/// tensor fields set to the tensor operation's tensor results.
+public protocol TensorGroup: TensorArrayProtocol {
+
+    /// The types of the tensor stored properties in this type.
+    static var _typeList: [TensorDataType] { get }
+
+    /// Initializes a value of this type, taking ownership of the `_tensorHandleCount` tensors
+    /// starting at address `tensorHandles`.
+    init(_owning tensorHandles: UnsafePointer<CTensorHandle>?)
+}
+
 public extension TensorGroup {
     /// The number of tensor fields in this type.
     static var _tensorHandleCount: Int32 { return Int32(Self._typeList.count) }