Re-implement tensor shape/rank operators using TFE runtime APIs. (#141)

Ported from swiftlang/swift#24949.

Gives a 98% speedup for `Tensor.rank` and 85% speedup for `Tensor.shape` (for a 2x3 tensor) on a MacBook Pro.

Author: eaplatanios

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

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