Add an API to get the underlying LazyTensorOperation of a LazyTensorHandle. (#405)

Author: bgogul

Sources/TensorFlow/Core/LazyTensorOperation.swift

@@ -73,6 +73,14 @@ class LazyTensorHandle: _AnyTensorHandle {
     /// The shape of the underlying `Tensor`.
     @inlinable
     var shape: TensorShape { _tfeTensorHandle.shape }
+
+    /// Returns the underlying `LazyTensorOperation` if this is a symbolic `LazyTensorHandle`.
+    var lazyTensorOperation: LazyTensorOperation? {
+        switch handle {
+        case .symbolic(let op, _, _): return op
+        case .concrete(_): return nil
+        }
+    }
 
     // Liveness tracking for LazyTensorOperations
     //

Tests/TensorFlowTests/LazyTensorHandleTests.swift

@@ -52,6 +52,21 @@ final class LazyTensorHandleTests: XCTestCase {
         XCTAssertEqual(liveSymTensor.description, "%0.0*")
     }
 
+    func testLazyTensorOperationProperty() {
+        let zero = Tensor<Float>(0.0)
+        let zeroTFEHandle = zero.handle.handle._tfeTensorHandle
+        let concTensor = LazyTensorHandle(zeroTFEHandle)
+        XCTAssertNil(concTensor.lazyTensorOperation)
+
+        let op = LazyTensorOperation(
+            _id: "0", name: "IdentityN", outputCount: 3)
+        let symTensor = LazyTensorHandle(_lazy: op, index: 0)
+        let lazyTensorOperation = symTensor.lazyTensorOperation
+        XCTAssertNotNil(lazyTensorOperation)
+        // Checks that returned value is the same as the one that we passed in.
+        XCTAssertTrue(lazyTensorOperation === op)
+    }
+
     func testLivenessTracking() {
         func assertLive(_ expectedLive: [LazyTensorOperation]) {
             var actualLiveOps: Set<LazyTensorOperationRef> = []
@@ -152,6 +167,7 @@ final class LazyTensorHandleTests: XCTestCase {
 
     static var allTests = [
         ("testConstructions", testConstructions),
+        ("testLazyTensorOperationProperty", testLazyTensorOperationProperty),
         ("testLivenessTracking", testLivenessTracking),
         ("testTensorToLazyTensorConversions", testTensorToLazyTensorConversions)
     ]