[TF] Re-implement tensor shape/rank operators using TFE runtime APIs.

Author: rxwei

stdlib/public/TensorFlow/CompilerRuntime.swift

@@ -554,7 +554,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()

stdlib/public/TensorFlow/Tensor.swift

@@ -367,7 +367,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)
     }
   }
 
@@ -376,14 +379,26 @@ 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 {
-    return Int(scalarCountTensor.scalar!)
+    @_semantics("autodiff.nonvarying")
+    get {
+      let status = _ExecutionContext.global.status
+      let size = TFE_TensorHandleNumElements(handle._cTensorHandle, status)
+      checkOk(status)
+      return Int(size)
+    }
   }
 }