Add liveness tracking tests.

Author: bgogul

Sources/TensorFlow/LazyTensor/LazyTensorOperation.swift

@@ -56,8 +56,8 @@ class LazyTensor: _AnyTensorHandle {
     //
     struct LazyTensorOperationRefCounts {
         let op: LazyTensorOperation
-        let live: Int
-        let all: Int
+        let liveRefCount: Int
+        let allRefCount: Int
     }
 
     private static var operationRefCounts: [
@@ -68,22 +68,22 @@ class LazyTensor: _AnyTensorHandle {
         if let counts = operationRefCounts[opId] {
             operationRefCounts[opId] = LazyTensorOperationRefCounts(
                 op: op,
-                live: isLive ? counts.live + 1 : counts.live,
-                all: counts.all + 1)
+                liveRefCount: counts.liveRefCount + (isLive ? 1 : 0),
+                allRefCount: counts.allRefCount + 1)
         } else {
             operationRefCounts[opId] = LazyTensorOperationRefCounts(
-                op: op, live: isLive ? 1 : 0, all: 1)
+                op: op, liveRefCount: isLive ? 1 : 0, allRefCount: 1)
         }
     }
 
     static func decrementRefCount(_ op: LazyTensorOperation, isLive: Bool) {
         let opId = ObjectIdentifier(op)
         if let counts = operationRefCounts[opId] {
-            if counts.all > 1 {
+            if counts.allRefCount > 1 {
                 operationRefCounts[opId] = LazyTensorOperationRefCounts(
                     op: op,
-                    live: isLive ? counts.live - 1 : counts.live,
-                    all: counts.all - 1)
+                    liveRefCount: counts.liveRefCount - (isLive ? 1 : 0),
+                    allRefCount: counts.allRefCount - 1)
             } else {
                 operationRefCounts.removeValue(forKey: opId)
             }
@@ -93,24 +93,20 @@ class LazyTensor: _AnyTensorHandle {
     static func isLive(_ op: LazyTensorOperation) -> Bool {
         let opId = ObjectIdentifier(op)
         if let counts = operationRefCounts[opId] {
-            return counts.live > 0
+            return counts.liveRefCount > 0
         }
         return false
     }
 
     static func onLiveOperations(_ perform: (LazyTensorOperation) -> ()) {
         for (_, counts) in operationRefCounts {
-            if (counts.live > 0) { perform(counts.op) }
+            if (counts.liveRefCount > 0) { perform(counts.op) }
         }
     }
 
-    static func onAllOperations(_ perform: (LazyTensorOperation) -> ()) {
-        for (_, counts) in operationRefCounts { perform(counts.op) }
-    }
-
-    public static func printRefCounts() {
+    static func printRefCounts() {
         let live = operationRefCounts.values.reduce(0, { (sum, element) in
-                return sum + (element.live > 0 ? 1 : 0)
+                return sum + (element.liveRefCount > 0 ? 1 : 0)
             })
         print("LazyTensorOperations: \(operationRefCounts.count) (\(live) live)")
     }

Tests/TensorFlowTests/LazyTensorTests.swift

@@ -16,6 +16,18 @@ import XCTest
 @testable import TensorFlow
 import CTensorFlow
 
+struct LazyTensorOperationRef : Equatable, Hashable {
+    let value: LazyTensorOperation
+    init(_ value: LazyTensorOperation) { self.value = value }
+    func hash(into hasher: inout Hasher) {
+        hasher.combine(ObjectIdentifier(value))
+    }
+}
+
+func ==(lhs: LazyTensorOperationRef, rhs: LazyTensorOperationRef) -> Bool {
+    return lhs.value === rhs.value
+}
+
 final class LazyTensorTests: XCTestCase {
     func testConstructions() {
         let zero = Tensor<Float>(0.0)
@@ -40,8 +52,54 @@ final class LazyTensorTests: XCTestCase {
         XCTAssertEqual("\(liveSymTensor)", "IdentityN_0:0*")
     }
 
+    func testLivenessTracking() {
+        func assertLive(_ expectedLive: [LazyTensorOperation]) {
+            var actualLiveOps: Set<LazyTensorOperationRef> = []
+            LazyTensor.onLiveOperations {
+                actualLiveOps.insert(LazyTensorOperationRef($0))
+            }
+            let expectedLiveOps = Set<LazyTensorOperationRef>(
+                expectedLive.map { LazyTensorOperationRef($0) }
+            )
+            XCTAssertEqual(expectedLiveOps, actualLiveOps)
+        }
+        func isSymbolic(_ t: LazyTensor) -> Bool {
+            if case let .symbolic(_) = t.handle {
+                return true
+            } else {
+                return false
+            }
+        }
+
+        let op0 = LazyTensorOperation(
+            _id: "0", name: "IdentityN", outputCount: 2)
+        let op1 = LazyTensorOperation(
+            _id: "1", name: "IdentityN", outputCount: 2)
+
+        XCTAssertFalse(LazyTensor.isLive(op0))
+        XCTAssertFalse(LazyTensor.isLive(op1))
+
+        let t0 = LazyTensor(_lazyLive: op0, index: 0)
+        let t1 = LazyTensor(_lazy: op1, index: 1)
+        XCTAssertTrue(LazyTensor.isLive(op0))
+        XCTAssertFalse(LazyTensor.isLive(op1))
+        do {
+            let t3 = LazyTensor(_lazyLive: op1, index: 0)
+            XCTAssertTrue(LazyTensor.isLive(op1))
+            assertLive([op0, op1])
+            // The following are here just to ensure t3 is live.
+            XCTAssertTrue(isSymbolic(t3))
+        }
+        XCTAssertFalse(LazyTensor.isLive(op1))
+        assertLive([op0])
+        // The following are here just to ensure t0 and t1 are live.
+        XCTAssertTrue(isSymbolic(t1))
+        XCTAssertTrue(isSymbolic(t0))
+    }
+
     static var allTests = [
         ("testConstructions", testConstructions),
+        ("testLivenessTracking", testLivenessTracking),
     ]
 
 }